1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
|
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])
|
