Query Testing

def django_assert_num_queries(pytestconfig: pytest.Config) -> DjangoAssertNumQueries:
    """
    Fixture for asserting exact database query counts.
    
    Returns a context manager that captures and asserts the exact number
    of database queries executed within its context. Useful for performance
    testing and ensuring optimal query patterns.
    
    Returns:
        DjangoAssertNumQueries: Query count assertion context manager
    """

class DjangoAssertNumQueries:
    """Context manager for asserting exact database query counts."""
    
    def __call__(
        self,
        num: int,
        connection: Any | None = None,
        info: str | None = None,
        *,
        using: str | None = None,
    ) -> ContextManager[None]:
        """
        Create context manager for exact/maximum query count assertion.
        
        Args:
            num: Expected exact/maximum number of database queries
            connection: Database connection to monitor (overrides using)
            info: Additional info to include in assertion error messages
            using: Database alias to monitor (default: "default")
            
        Returns:
            ContextManager: Context manager that asserts query count
        """

def test_query_count(django_assert_num_queries):
    from myapp.models import MyModel
    
    # Assert exactly 1 query (the INSERT)
    with django_assert_num_queries(1):
        MyModel.objects.create(name="test")
    
    # Assert exactly 2 queries (SELECT and UPDATE)
    with django_assert_num_queries(2):
        obj = MyModel.objects.get(name="test")
        obj.name = "updated"
        obj.save()

def django_assert_max_num_queries(pytestconfig: pytest.Config) -> DjangoAssertNumQueries:
    """
    Fixture for asserting maximum database query counts.
    
    Returns a context manager that captures database queries and asserts
    that the count does not exceed the specified maximum. Useful for
    performance regression testing.
    
    Returns:
        DjangoAssertNumQueries: Maximum query count assertion context manager
    """

def test_max_query_count(django_assert_max_num_queries):
    from myapp.models import MyModel, RelatedModel
    
    # Ensure operation uses at most 3 queries
    with django_assert_max_num_queries(3):
        # This should be optimized to avoid N+1 queries
        for obj in MyModel.objects.select_related('related'):
            print(obj.related.name)
    
    # Performance regression test
    with django_assert_max_num_queries(5):
        # Complex operation that should stay under 5 queries
        result = MyModel.objects.filter(active=True).prefetch_related('tags')
        for obj in result:
            obj.tags.all()

from typing import ContextManager, List, Optional
from django.db import connection

# Query assertion context managers
class DjangoAssertNumQueries:
    """Context manager for exact query count assertions."""
    
    def __call__(self, num: int, using: str = "default") -> ContextManager[None]:
        """
        Assert exact number of queries.
        
        Args:
            num: Expected exact number of queries
            using: Database alias to monitor (default: "default")
            
        Returns:
            ContextManager that captures and asserts query count
        """


# Query information types
class QueryInfo:
    """Information about executed database query."""
    sql: str
    time: float
    
# Database connection query capture
class ConnectionQueryCapture:
    """Capture queries from database connection."""
    queries: List[QueryInfo]
    initial_queries: int
    final_queries: int
    executed_queries: List[QueryInfo]