wrapper to avoid accessing connection attributes until inside coroutine function connection access is thread sensitive and cannot be passed across syncasync boundaries evaluate the queryset to populate the query cache the query cache is used unsaved instances are not allowed so use an id known not to exist todo check the captured query when async versions of self assertnumqueriescapturequeriescontext context processors are available wrapper to avoid accessing connection attributes until inside coroutine function connection access is thread sensitive and cannot be passed across sync async boundaries evaluate the queryset to populate the query cache the query cache is used unsaved instances are not allowed so use an id known not to exist todo check the captured query when async versions of self assertnumqueries capturequeriescontext context processors are available
import json
import xml.etree.ElementTree
from datetime import datetime

from asgiref.sync import async_to_sync, sync_to_async

from django.db import NotSupportedError, connection
from django.db.models import Prefetch, Sum
from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature

from .models import RelatedModel, SimpleModel


class AsyncQuerySetTest(TestCase):
    @classmethod
    def setUpTestData(cls):
        cls.s1 = SimpleModel.objects.create(
            field=1,
            created=datetime(2022, 1, 1, 0, 0, 0),
        )
        cls.s2 = SimpleModel.objects.create(
            field=2,
            created=datetime(2022, 1, 1, 0, 0, 1),
        )
        cls.s3 = SimpleModel.objects.create(
            field=3,
            created=datetime(2022, 1, 1, 0, 0, 2),
        )
        cls.r1 = RelatedModel.objects.create(simple=cls.s1)
        cls.r2 = RelatedModel.objects.create(simple=cls.s2)
        cls.r3 = RelatedModel.objects.create(simple=cls.s3)

    @staticmethod
    def _get_db_feature(connection_, feature_name):
        
        
        
        return getattr(connection_.features, feature_name)

    async def test_async_iteration(self):
        results = []
        async for m in SimpleModel.objects.order_by("pk"):
            results.append(m)
        self.assertEqual(results, [self.s1, self.s2, self.s3])

    async def test_aiterator(self):
        qs = SimpleModel.objects.aiterator()
        results = []
        async for m in qs:
            results.append(m)
        self.assertCountEqual(results, [self.s1, self.s2, self.s3])

    async def test_aiterator_prefetch_related(self):
        results = []
        async for s in SimpleModel.objects.prefetch_related(
            Prefetch("relatedmodel_set", to_attr="prefetched_relatedmodel")
        ).aiterator():
            results.append(s.prefetched_relatedmodel)
        self.assertCountEqual(results, [[self.r1], [self.r2], [self.r3]])

    async def test_aiterator_invalid_chunk_size(self):
        msg = "Chunk size must be strictly positive."
        for size in [0, -1]:
            qs = SimpleModel.objects.aiterator(chunk_size=size)
            with self.subTest(size=size), self.assertRaisesMessage(ValueError, msg):
                async for m in qs:
                    pass

    async def test_acount(self):
        count = await SimpleModel.objects.acount()
        self.assertEqual(count, 3)

    async def test_acount_cached_result(self):
        qs = SimpleModel.objects.all()
        
        [x async for x in qs]
        count = await qs.acount()
        self.assertEqual(count, 3)

        await sync_to_async(SimpleModel.objects.create)(
            field=4,
            created=datetime(2022, 1, 1, 0, 0, 0),
        )
        
        count = await qs.acount()
        self.assertEqual(count, 3)

    async def test_aget(self):
        instance = await SimpleModel.objects.aget(field=1)
        self.assertEqual(instance, self.s1)

    async def test_acreate(self):
        await SimpleModel.objects.acreate(field=4)
        self.assertEqual(await SimpleModel.objects.acount(), 4)

    async def test_aget_or_create(self):
        instance, created = await SimpleModel.objects.aget_or_create(field=4)
        self.assertEqual(await SimpleModel.objects.acount(), 4)
        self.assertIs(created, True)

    async def test_aupdate_or_create(self):
        instance, created = await SimpleModel.objects.aupdate_or_create(
            id=self.s1.id, defaults={"field": 2}
        )
        self.assertEqual(instance, self.s1)
        self.assertEqual(instance.field, 2)
        self.assertIs(created, False)
        instance, created = await SimpleModel.objects.aupdate_or_create(field=4)
        self.assertEqual(await SimpleModel.objects.acount(), 4)
        self.assertIs(created, True)
        instance, created = await SimpleModel.objects.aupdate_or_create(
            field=5, defaults={"field": 7}, create_defaults={"field": 6}
        )
        self.assertEqual(await SimpleModel.objects.acount(), 5)
        self.assertIs(created, True)
        self.assertEqual(instance.field, 6)

    @skipUnlessDBFeature("has_bulk_insert")
    @async_to_sync
    async def test_abulk_create(self):
        instances = [SimpleModel(field=i) for i in range(10)]
        qs = await SimpleModel.objects.abulk_create(instances)
        self.assertEqual(len(qs), 10)

    @skipUnlessDBFeature("has_bulk_insert", "supports_update_conflicts")
    @skipIfDBFeature("supports_update_conflicts_with_target")
    @async_to_sync
    async def test_update_conflicts_unique_field_unsupported(self):
        msg = (
            "This database backend does not support updating conflicts with specifying "
            "unique fields that can trigger the upsert."
        )
        with self.assertRaisesMessage(NotSupportedError, msg):
            await SimpleModel.objects.abulk_create(
                [SimpleModel(field=1), SimpleModel(field=2)],
                update_conflicts=True,
                update_fields=["field"],
                unique_fields=["created"],
            )

    async def test_abulk_update(self):
        instances = SimpleModel.objects.all()
        async for instance in instances:
            instance.field = instance.field * 10

        await SimpleModel.objects.abulk_update(instances, ["field"])

        qs = [(o.pk, o.field) async for o in SimpleModel.objects.all()]
        self.assertCountEqual(
            qs,
            [(self.s1.pk, 10), (self.s2.pk, 20), (self.s3.pk, 30)],
        )

    async def test_ain_bulk(self):
        res = await SimpleModel.objects.ain_bulk()
        self.assertEqual(
            res,
            {self.s1.pk: self.s1, self.s2.pk: self.s2, self.s3.pk: self.s3},
        )

        res = await SimpleModel.objects.ain_bulk([self.s2.pk])
        self.assertEqual(res, {self.s2.pk: self.s2})

        res = await SimpleModel.objects.ain_bulk([self.s2.pk], field_name="id")
        self.assertEqual(res, {self.s2.pk: self.s2})

    async def test_alatest(self):
        instance = await SimpleModel.objects.alatest("created")
        self.assertEqual(instance, self.s3)

        instance = await SimpleModel.objects.alatest("-created")
        self.assertEqual(instance, self.s1)

    async def test_aearliest(self):
        instance = await SimpleModel.objects.aearliest("created")
        self.assertEqual(instance, self.s1)

        instance = await SimpleModel.objects.aearliest("-created")
        self.assertEqual(instance, self.s3)

    async def test_afirst(self):
        instance = await SimpleModel.objects.afirst()
        self.assertEqual(instance, self.s1)

        instance = await SimpleModel.objects.filter(field=4).afirst()
        self.assertIsNone(instance)

    async def test_alast(self):
        instance = await SimpleModel.objects.alast()
        self.assertEqual(instance, self.s3)

        instance = await SimpleModel.objects.filter(field=4).alast()
        self.assertIsNone(instance)

    async def test_aaggregate(self):
        total = await SimpleModel.objects.aaggregate(total=Sum("field"))
        self.assertEqual(total, {"total": 6})

    async def test_aexists(self):
        check = await SimpleModel.objects.filter(field=1).aexists()
        self.assertIs(check, True)

        check = await SimpleModel.objects.filter(field=4).aexists()
        self.assertIs(check, False)

    async def test_acontains(self):
        check = await SimpleModel.objects.acontains(self.s1)
        self.assertIs(check, True)
        
        check = await SimpleModel.objects.acontains(
            SimpleModel(id=self.s3.id + 1, field=4)
        )
        self.assertIs(check, False)

    async def test_aupdate(self):
        await SimpleModel.objects.aupdate(field=99)
        qs = [o async for o in SimpleModel.objects.all()]
        values = [instance.field for instance in qs]
        self.assertEqual(set(values), {99})

    async def test_adelete(self):
        await SimpleModel.objects.filter(field=2).adelete()
        qs = [o async for o in SimpleModel.objects.all()]
        self.assertCountEqual(qs, [self.s1, self.s3])

    @skipUnlessDBFeature("supports_explaining_query_execution")
    @async_to_sync
    async def test_aexplain(self):
        supported_formats = await sync_to_async(self._get_db_feature)(
            connection, "supported_explain_formats"
        )
        all_formats = (None, *supported_formats)
        for format_ in all_formats:
            with self.subTest(format=format_):
                
                
                
                result = await SimpleModel.objects.filter(field=1).aexplain(
                    format=format_
                )
                self.assertIsInstance(result, str)
                self.assertTrue(result)
                if not format_:
                    continue
                if format_.lower() == "xml":
                    try:
                        xml.etree.ElementTree.fromstring(result)
                    except xml.etree.ElementTree.ParseError as e:
                        self.fail(f"QuerySet.aexplain() result is not valid XML: {e}")
                elif format_.lower() == "json":
                    try:
                        json.loads(result)
                    except json.JSONDecodeError as e:
                        self.fail(f"QuerySet.aexplain() result is not valid JSON: {e}")

    async def test_raw(self):
        sql = "SELECT id, field FROM async_simplemodel WHERE created=%s"
        qs = SimpleModel.objects.raw(sql, [self.s1.created])
        self.assertEqual([o async for o in qs], [self.s1])