22
33from django .core .exceptions import FieldDoesNotExist
44from django .db import connection , models
5+ from django .db .models .expressions import Value
56from django .test import SimpleTestCase , TestCase
67from django .test .utils import CaptureQueriesContext , isolate_apps
78
8- from django_mongodb_backend .fields import EmbeddedModelArrayField
9+ from django_mongodb_backend .fields import ArrayField , EmbeddedModelArrayField
910from django_mongodb_backend .models import EmbeddedModel
1011
1112from .models import Artifact , Audit , Exhibit , Movie , Restoration , Review , Section , Tour
@@ -147,9 +148,9 @@ def setUpTestData(cls):
147148 cls .egypt_tour = Tour .objects .create (guide = "Amira" , exhibit = cls .egypt )
148149 cls .wonders_tour = Tour .objects .create (guide = "Carlos" , exhibit = cls .wonders )
149150 cls .lost_tour = Tour .objects .create (guide = "Yelena" , exhibit = cls .lost_empires )
150- cls .audit_1 = Audit .objects .create (related_section_number = 1 , reviewed = True )
151- cls .audit_2 = Audit .objects .create (related_section_number = 2 , reviewed = True )
152- cls .audit_3 = Audit .objects .create (related_section_number = 5 , reviewed = False )
151+ cls .audit_1 = Audit .objects .create (section_number = 1 , reviewed = True )
152+ cls .audit_2 = Audit .objects .create (section_number = 2 , reviewed = True )
153+ cls .audit_3 = Audit .objects .create (section_number = 5 , reviewed = False )
153154
154155 def test_exact (self ):
155156 self .assertCountEqual (
@@ -288,34 +289,60 @@ def test_foreign_field_with_slice(self):
288289 qs = Tour .objects .filter (exhibit__sections__0_2__section_number__in = [1 , 2 ])
289290 self .assertCountEqual (qs , [self .wonders_tour , self .egypt_tour ])
290291
291- def test_subquery_section_number_lt (self ):
292+ def test_subquery_numeric_lookups (self ):
292293 subq = Audit .objects .filter (
293- related_section_number__in = models .OuterRef ("sections__section_number" )
294- ).values ("related_section_number" )[:1 ]
295- self .assertCountEqual (
296- Exhibit .objects .filter (sections__section_number = subq ),
297- [self .egypt , self .wonders , self .new_descoveries ],
298- )
294+ section_number__in = models .OuterRef ("sections__section_number" )
295+ ).values ("section_number" )[:1 ]
296+ lookups_expected_value = [
297+ ("exact" , [self .egypt , self .new_descoveries , self .wonders ]),
298+ ("lt" , []),
299+ ("lte" , [self .egypt , self .new_descoveries , self .wonders ]),
300+ ("gt" , [self .wonders ]),
301+ ("gte" , [self .egypt , self .new_descoveries , self .wonders ]),
302+ ]
303+ for lookup , expected_value in lookups_expected_value :
304+ with self .subTest (key = lookup ):
305+ kwargs = {f"sections__section_number__{ lookup } " : subq }
306+ self .assertCountEqual (Exhibit .objects .filter (** kwargs ), expected_value )
299307
300- def test_check_in_subquery (self ):
301- subquery = Audit .objects .filter (reviewed = True ).values_list (
302- "related_section_number" , flat = True
303- )
308+ def test_in_subquery (self ):
309+ subquery = Audit .objects .filter (reviewed = True ).values_list ("section_number" , flat = True )
304310 result = Exhibit .objects .filter (sections__section_number__in = subquery )
305- self .assertCountEqual (result , [self .wonders , self .egypt , self .new_descoveries ])
311+ self .assertCountEqual (result , [self .wonders , self .new_descoveries , self .egypt ])
306312
307313 def test_array_as_rhs (self ):
308314 result = Exhibit .objects .filter (
309315 main_section__section_number__in = models .F ("sections__section_number" )
310316 )
311317 self .assertCountEqual (result , [self .new_descoveries ])
312318
313- def test_array_annotation_as_rhs (self ):
319+ def test_array_annotation_lookup (self ):
314320 result = Exhibit .objects .annotate (
315321 section_numbers = models .F ("main_section__section_number" )
316322 ).filter (section_numbers__in = models .F ("sections__section_number" ))
317323 self .assertCountEqual (result , [self .new_descoveries ])
318324
325+ def test_arrayfield_annotation_lookups (self ):
326+ lookup_expected_value = [
327+ ("exact" , [self .wonders ]),
328+ ("lt" , [self .new_descoveries ]),
329+ ("lte" , [self .wonders , self .new_descoveries ]),
330+ ("gt" , [self .egypt , self .lost_empires ]),
331+ ("gte" , [self .egypt , self .wonders , self .lost_empires ]),
332+ ("overlap" , [self .egypt , self .wonders , self .new_descoveries ]),
333+ ("contained_by" , [self .wonders ]),
334+ ("contains" , [self .egypt , self .wonders , self .new_descoveries , self .lost_empires ]),
335+ ]
336+ for lookup , expected in lookup_expected_value :
337+ with self .subTest (key = lookup ):
338+ kwargs = {f"section_numbers__{ lookup } " : models .F ("sections__section_number" )}
339+ result = Exhibit .objects .annotate (
340+ section_numbers = Value (
341+ [1 , 2 ], output_field = ArrayField (base_field = models .IntegerField ())
342+ )
343+ ).filter (** kwargs )
344+ self .assertCountEqual (result , expected )
345+
319346 def test_array_annotation_index (self ):
320347 # This test would be useful to have, but it cannot be implemented
321348 # due to the current annotation handling.
@@ -327,9 +354,7 @@ def test_array_annotation_index(self):
327354 self .assertCountEqual (result , [self .new_descoveries , self .egypt ])
328355
329356 def test_array_annotation (self ):
330- results = Exhibit .objects .annotate (
331- section_numbers = models .F ("sections__section_number" )
332- ).all ()
357+ results = Exhibit .objects .annotate (section_numbers = models .F ("sections__section_number" ))
333358 sections_numbers = [e .section_numbers for e in results ]
334359 self .assertCountEqual (sections_numbers , [[1 ], [1 , 2 ], [2 ], []])
335360
0 commit comments