django-rest-framework/tests/test_optimization.py

"""
Tests for query optimization features in Django REST Framework.
"""

from django.contrib.auth.models import Group, User
from django.test import TestCase

from rest_framework import generics, serializers, viewsets
from rest_framework.optimization import (
    OptimizedQuerySetMixin,
    analyze_serializer_fields,
    detect_n_plus_one,
    get_optimization_suggestions,
    optimize_queryset,
)
from rest_framework.test import APIRequestFactory

from .models import (
    ForeignKeyTarget, ForeignKeySource, ManyToManyTarget, ManyToManySource
)

factory = APIRequestFactory()


# Test serializers using existing models
class ForeignKeyTargetSerializer(serializers.ModelSerializer):
    class Meta:
        model = ForeignKeyTarget
        fields = '__all__'


class ForeignKeySourceSerializer(serializers.ModelSerializer):
    target = ForeignKeyTargetSerializer(read_only=True)

    class Meta:
        model = ForeignKeySource
        fields = '__all__'


class ForeignKeySourceListSerializer(serializers.ModelSerializer):
    """Simpler serializer for list view."""
    class Meta:
        model = ForeignKeySource
        fields = ['id', 'name', 'target']


class ManyToManySourceSerializer(serializers.ModelSerializer):
    targets = serializers.PrimaryKeyRelatedField(many=True, read_only=True)

    class Meta:
        model = ManyToManySource
        fields = '__all__'


class TestQueryAnalyzer(TestCase):
    """Test the QueryAnalyzer functionality."""

    def test_analyze_serializer_with_foreign_key(self):
        """Test analyzing a serializer with ForeignKey relationships."""
        analysis = analyze_serializer_fields(ForeignKeySourceListSerializer)
        
        self.assertIn('select_related', analysis)
        self.assertIn('prefetch_related', analysis)
        # Should detect target as select_related
        self.assertIn('target', analysis['select_related'])

    def test_analyze_serializer_with_many_to_many(self):
        """Test analyzing a serializer with ManyToMany relationships."""
        analysis = analyze_serializer_fields(ManyToManySourceSerializer)
        
        # Should detect targets as prefetch_related
        self.assertIn('targets', analysis['prefetch_related'])

    def test_analyze_nested_serializer(self):
        """Test analyzing nested serializers."""
        analysis = analyze_serializer_fields(ForeignKeySourceSerializer)
        
        # Should detect target as select_related
        self.assertIn('target', analysis['select_related'])

    def test_analyze_non_model_serializer(self):
        """Test analyzing a non-ModelSerializer returns empty analysis."""
        class SimpleSerializer(serializers.Serializer):
            name = serializers.CharField()
        
        analysis = analyze_serializer_fields(SimpleSerializer)
        
        self.assertEqual(analysis['select_related'], [])
        self.assertEqual(analysis['prefetch_related'], [])


class TestQueryOptimizer(TestCase):
    """Test the query optimizer functionality."""

    def setUp(self):
        self.target = ForeignKeyTarget.objects.create(name='Test Target')
        self.source = ForeignKeySource.objects.create(name='Test Source', target=self.target)
        self.m2m_target = ManyToManyTarget.objects.create(name='M2M Target')
        self.m2m_source = ManyToManySource.objects.create(name='M2M Source')
        self.m2m_source.targets.add(self.m2m_target)

    def test_optimize_queryset_with_select_related(self):
        """Test optimizing a queryset with select_related."""
        queryset = ForeignKeySource.objects.all()
        optimized = optimize_queryset(queryset, ForeignKeySourceListSerializer, auto_optimize=True)
        
        # Check that select_related was applied
        self.assertTrue(hasattr(optimized.query, 'select_related'))
        # The queryset should have select_related applied
        self.assertTrue(optimized.query.select_related)

    def test_optimize_queryset_with_prefetch_related(self):
        """Test optimizing a queryset with prefetch_related."""
        queryset = ManyToManySource.objects.all()
        optimized = optimize_queryset(queryset, ManyToManySourceSerializer, auto_optimize=True)
        
        # Check that prefetch_related was applied
        self.assertTrue(hasattr(optimized.query, 'prefetch_related_lookups'))
        self.assertIn('targets', optimized.query.prefetch_related_lookups)

    def test_optimize_queryset_explicit_fields(self):
        """Test optimizing with explicitly provided fields."""
        queryset = ForeignKeySource.objects.all()
        optimized = optimize_queryset(
            queryset,
            ForeignKeySourceListSerializer,
            select_related=['target'],
            auto_optimize=False
        )
        
        self.assertTrue(optimized.query.select_related)

    def test_optimize_queryset_merges_explicit_and_auto(self):
        """Test that explicit and auto-detected fields are merged."""
        queryset = ForeignKeySource.objects.all()
        optimized = optimize_queryset(
            queryset,
            ForeignKeySourceListSerializer,
            select_related=['target'],
            auto_optimize=True
        )
        
        # Should have both explicit and auto-detected fields
        self.assertTrue(optimized.query.select_related)

    def test_optimize_queryset_non_queryset(self):
        """Test that non-queryset objects are returned as-is."""
        result = optimize_queryset([1, 2, 3], ForeignKeySourceListSerializer)
        self.assertEqual(result, [1, 2, 3])


class TestOptimizedQuerySetMixin(TestCase):
    """Test the OptimizedQuerySetMixin."""

    def setUp(self):
        self.target = ForeignKeyTarget.objects.create(name='Test Target')
        self.source = ForeignKeySource.objects.create(name='Test Source', target=self.target)
        self.m2m_target = ManyToManyTarget.objects.create(name='M2M Target')
        self.m2m_source = ManyToManySource.objects.create(name='M2M Source')
        self.m2m_source.targets.add(self.m2m_target)

    def test_mixin_applies_optimization(self):
        """Test that the mixin applies optimizations automatically."""
        class SourceViewSet(OptimizedQuerySetMixin, viewsets.ReadOnlyModelViewSet):
            queryset = ForeignKeySource.objects.all()
            serializer_class = ForeignKeySourceListSerializer
        
        view = SourceViewSet()
        queryset = view.get_queryset()
        
        # Should have optimizations applied
        self.assertTrue(hasattr(queryset.query, 'select_related'))

    def test_mixin_with_explicit_fields(self):
        """Test mixin with explicitly specified fields."""
        class SourceViewSet(OptimizedQuerySetMixin, viewsets.ReadOnlyModelViewSet):
            queryset = ForeignKeySource.objects.all()
            serializer_class = ForeignKeySourceListSerializer
            select_related_fields = ['target']
        
        view = SourceViewSet()
        queryset = view.get_queryset()
        
        self.assertTrue(queryset.query.select_related)


class TestDetectNPlusOne(TestCase):
    """Test N+1 query detection."""

    def setUp(self):
        self.target = ForeignKeyTarget.objects.create(name='Test Target')
        self.source = ForeignKeySource.objects.create(name='Test Source', target=self.target)

    def test_detect_n_plus_one_without_optimization(self):
        """Test detecting N+1 queries when queryset is not optimized."""
        queryset = ForeignKeySource.objects.all()
        warnings = detect_n_plus_one(ForeignKeySourceListSerializer, queryset)
        
        # Should detect potential N+1 queries
        self.assertGreater(len(warnings), 0)
        self.assertTrue(any('target' in w for w in warnings))

    def test_detect_n_plus_one_with_optimization(self):
        """Test that optimized querysets don't trigger warnings."""
        queryset = ForeignKeySource.objects.select_related('target')
        warnings = detect_n_plus_one(ForeignKeySourceListSerializer, queryset)
        
        # Should have fewer or no warnings for optimized queryset
        # (Note: This may need adjustment based on implementation)

    def test_detect_n_plus_one_non_queryset(self):
        """Test that non-queryset objects return empty warnings."""
        warnings = detect_n_plus_one(ForeignKeySourceListSerializer, [1, 2, 3])
        self.assertEqual(warnings, [])


class TestOptimizationSuggestions(TestCase):
    """Test getting optimization suggestions."""

    def test_get_suggestions(self):
        """Test getting optimization suggestions for a serializer."""
        suggestions = get_optimization_suggestions(ForeignKeySourceListSerializer)
        
        self.assertIn('select_related', suggestions)
        self.assertIn('prefetch_related', suggestions)
        self.assertIn('code_example', suggestions)
        
        # Should have suggestions
        self.assertGreater(len(suggestions['select_related']), 0)
        
        # Should have code example
        if suggestions['select_related'] or suggestions['prefetch_related']:
            self.assertIsNotNone(suggestions['code_example'])
            self.assertIn('get_queryset', suggestions['code_example'])