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): # Wrapper to avoid accessing connection attributes until inside # coroutine function. Connection access is thread sensitive and cannot # be passed across sync/async boundaries. 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() # Evaluate the queryset to populate the query cache. [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), ) # The query cache is used. 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) # Unsaved instances are not allowed, so use an ID known not to exist. 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_): # TODO: Check the captured query when async versions of # self.assertNumQueries/CaptureQueriesContext context # processors are available. 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])