33
44from django_mongodb_backend .indexes import SearchIndex , VectorSearchIndex
55
6- from .models import Article
6+ from .models import SearchIndexTestModel
77
88
99@skipIfDBFeature ("supports_atlas_search" )
1010class UnsupportedSearchIndexesTests (TestCase ):
1111 def test_search_index_not_created (self ):
1212 with connection .schema_editor () as editor :
13- index = SearchIndex (name = "recent_article_idx " , fields = ["number" ])
13+ index = SearchIndex (name = "recent_test_idx " , fields = ["number" ])
1414 with self .assertNumQueries (0 ):
15- editor .add_index (index = index , model = Article )
16- self .assertNotIn (
17- index .name ,
18- connection .introspection .get_constraints (
19- cursor = None ,
20- table_name = Article ._meta .db_table ,
21- ),
22- )
15+ editor .add_index (index = index , model = SearchIndexTestModel )
16+ self .assertNotIn (
17+ index .name ,
18+ connection .introspection .get_constraints (
19+ cursor = None ,
20+ table_name = SearchIndexTestModel ._meta .db_table ,
21+ ),
22+ )
2323
2424 def test_vector_index_not_created (self ):
2525 with connection .schema_editor () as editor :
26- index = VectorSearchIndex (name = "recent_article_idx " , fields = ["number" ])
26+ index = VectorSearchIndex (name = "recent_test_idx " , fields = ["number" ])
2727 with self .assertNumQueries (0 ):
28- editor .add_index (index = index , model = Article )
29- self .assertNotIn (
30- index .name ,
31- connection .introspection .get_constraints (
32- cursor = None ,
33- table_name = Article ._meta .db_table ,
34- ),
35- )
28+ editor .add_index (index = index , model = SearchIndexTestModel )
29+ self .assertNotIn (
30+ index .name ,
31+ connection .introspection .get_constraints (
32+ cursor = None ,
33+ table_name = SearchIndexTestModel ._meta .db_table ,
34+ ),
35+ )
3636
3737
3838class VectorSearchIndexTests (SimpleTestCase ):
3939 def test_deconstruct_default_similarity (self ):
40- index = VectorSearchIndex (name = "recent_article_idx " , fields = ["number" ])
40+ index = VectorSearchIndex (name = "recent_test_idx " , fields = ["number" ])
4141 name , args , kwargs = index .deconstruct ()
4242 new = VectorSearchIndex (* args , ** kwargs )
4343 self .assertEqual (new .similarities , index .similarities )
4444
4545 def test_deconstruct_with_similarities (self ):
4646 index = VectorSearchIndex (
47- name = "recent_article_idx " ,
48- fields = ["number" , "headline " ],
47+ name = "recent_test_idx " ,
48+ fields = ["number" , "text " ],
4949 similarities = ["cosine" , "dotProduct" ],
5050 )
5151 path , args , kwargs = index .deconstruct ()
5252 self .assertEqual (
5353 kwargs ,
5454 {
55- "name" : "recent_article_idx " ,
56- "fields" : ["number" , "headline " ],
55+ "name" : "recent_test_idx " ,
56+ "fields" : ["number" , "text " ],
5757 "similarities" : ["cosine" , "dotProduct" ],
5858 },
5959 )
@@ -79,110 +79,145 @@ def test_define_field_twice(self):
7979 )
8080
8181
82+ class SchemaAssertionMixin :
83+ def assertExistAndRemoveIndex (self , model , index ):
84+ self .assertIn (
85+ index .name ,
86+ connection .introspection .get_constraints (
87+ cursor = None ,
88+ table_name = model ._meta .db_table ,
89+ ),
90+ )
91+ with connection .schema_editor () as editor :
92+ editor .remove_index (index = index , model = model )
93+ self .assertNotIn (
94+ index .name ,
95+ connection .introspection .get_constraints (
96+ cursor = None ,
97+ table_name = model ._meta .db_table ,
98+ ),
99+ )
100+
101+
82102@skipUnlessDBFeature ("supports_atlas_search" )
83- class SearchIndexSchemaTests (TestCase ):
103+ class SearchIndexSchemaTests (SchemaAssertionMixin , TestCase ):
84104 def test_simple (self ):
105+ index = SearchIndex (
106+ name = "recent_test_idx" ,
107+ fields = ["text" ],
108+ )
85109 with connection .schema_editor () as editor :
86- index = SearchIndex (
87- name = "recent_article_idx" ,
88- fields = ["number" ],
89- )
90- editor .add_index (index = index , model = Article )
91- editor .remove_index (index = index , model = Article )
92-
93- def test_multiple_fields (self ):
110+ editor .add_index (index = index , model = SearchIndexTestModel )
111+ self .assertExistAndRemoveIndex (index = index , model = SearchIndexTestModel )
112+
113+ def test_all_fields (self ):
114+ index = SearchIndex (
115+ name = "recent_test_idx" ,
116+ fields = [
117+ "text" ,
118+ "object_id" ,
119+ "number" ,
120+ "embedded" ,
121+ "json_data" ,
122+ "vector_integer" ,
123+ "vector_float" ,
124+ "boolean" ,
125+ "date" ,
126+ ],
127+ )
94128 with connection .schema_editor () as editor :
95- index = SearchIndex (
96- name = "recent_article_idx" ,
97- fields = [
98- "headline" ,
99- "number" ,
100- "body" ,
101- "data" ,
102- "embedded" ,
103- "created_at" ,
104- "object_id" ,
105- "genres" ,
106- ],
107- )
108- editor .add_index (index = index , model = Article )
109- index_info = connection .introspection .get_constraints (
110- cursor = None ,
111- table_name = Article ._meta .db_table ,
112- )
113- expected_options = {
114- "dynamic" : False ,
115- "fields" : {
116- "created_at" : {"type" : "date" },
117- "body" : {
118- "indexOptions" : "offsets" ,
119- "norms" : "include" ,
120- "store" : True ,
121- "type" : "string" ,
122- },
123- "data" : {"dynamic" : False , "fields" : {}, "type" : "document" },
124- "embedded" : {"dynamic" : False , "fields" : {}, "type" : "embeddedDocuments" },
125- "headline" : {
126- "indexOptions" : "offsets" ,
127- "norms" : "include" ,
128- "store" : True ,
129- "type" : "string" ,
130- },
131- "number" : {
132- "indexDoubles" : True ,
133- "indexIntegers" : True ,
134- "representation" : "double" ,
135- "type" : "number" ,
136- },
137- "object_id" : {"type" : "objectId" },
138- "genres" : {"dynamic" : False , "fields" : {}, "type" : "embeddedDocuments" },
129+ editor .add_index (index = index , model = SearchIndexTestModel )
130+ index_info = connection .introspection .get_constraints (
131+ cursor = None ,
132+ table_name = SearchIndexTestModel ._meta .db_table ,
133+ )
134+ expected_options = {
135+ "dynamic" : False ,
136+ "fields" : {
137+ "boolean" : {"type" : "boolean" },
138+ "embedded" : {"dynamic" : False , "fields" : {}, "type" : "embeddedDocuments" },
139+ "json_data" : {"dynamic" : False , "fields" : {}, "type" : "document" },
140+ "number" : {
141+ "indexDoubles" : True ,
142+ "indexIntegers" : True ,
143+ "representation" : "double" ,
144+ "type" : "number" ,
145+ },
146+ "object_id" : {"type" : "objectId" },
147+ "text" : {
148+ "indexOptions" : "offsets" ,
149+ "norms" : "include" ,
150+ "store" : True ,
151+ "type" : "string" ,
139152 },
140- }
141- self .assertCountEqual (index_info [index .name ]["columns" ], index .fields )
142- self .assertEqual (index_info [index .name ]["options" ], expected_options )
143- editor .remove_index (index = index , model = Article )
153+ "date" : {"type" : "date" },
154+ "vector_float" : {"dynamic" : False , "fields" : {}, "type" : "embeddedDocuments" },
155+ "vector_integer" : {"dynamic" : False , "fields" : {}, "type" : "embeddedDocuments" },
156+ },
157+ }
158+ self .assertCountEqual (index_info [index .name ]["columns" ], index .fields )
159+ self .assertEqual (index_info [index .name ]["options" ], expected_options )
160+ self .assertExistAndRemoveIndex (index = index , model = SearchIndexTestModel )
144161
145162
146163@skipUnlessDBFeature ("supports_atlas_search" )
147- class VectorSearchIndexSchemaTests (TestCase ):
164+ class VectorSearchIndexSchemaTests (SchemaAssertionMixin , TestCase ):
148165 def test_simple_vector_search (self ):
166+ index = VectorSearchIndex (name = "recent_test_idx" , fields = ["number" ])
149167 with connection .schema_editor () as editor :
150- index = VectorSearchIndex (name = "recent_article_idx" , fields = ["number" ])
151- editor .add_index (index = index , model = Article )
152- editor .remove_index (index = index , model = Article )
168+ editor .add_index (index = index , model = SearchIndexTestModel )
169+ self .assertExistAndRemoveIndex (index = index , model = SearchIndexTestModel )
153170
154171 def test_multiple_fields (self ):
172+ index = VectorSearchIndex (
173+ name = "recent_test_idx" ,
174+ fields = [
175+ "text" ,
176+ "object_id" ,
177+ "number" ,
178+ "embedded" ,
179+ "vector_integer" ,
180+ "vector_float" ,
181+ "boolean" ,
182+ "date" ,
183+ ],
184+ )
155185 with connection .schema_editor () as editor :
156- index = VectorSearchIndex (
157- name = "recent_article_idx" ,
158- fields = ["headline" , "number" , "body" , "description_semantic" ],
159- )
160- editor .add_index (index = index , model = Article )
161- index_info = connection .introspection .get_constraints (
162- cursor = None ,
163- table_name = Article ._meta .db_table ,
164- )
165- expected_options = {
166- "latestDefinition" : {
167- "fields" : [
168- {"path" : "headline" , "type" : "filter" },
169- {"path" : "number" , "type" : "filter" },
170- {"path" : "body" , "type" : "filter" },
171- {
172- "numDimensions" : 10 ,
173- "path" : "description_semantic" ,
174- "similarity" : "cosine" ,
175- "type" : "vector" ,
176- },
177- ]
178- },
179- "latestVersion" : 0 ,
180- "name" : "recent_article_idx" ,
181- "queryable" : False ,
182- "type" : "vectorSearch" ,
183- }
184- self .assertCountEqual (index_info [index .name ]["columns" ], index .fields )
185- index_info [index .name ]["options" ].pop ("id" )
186- index_info [index .name ]["options" ].pop ("status" )
187- self .assertEqual (index_info [index .name ]["options" ], expected_options )
188- editor .remove_index (index = index , model = Article )
186+ editor .add_index (index = index , model = SearchIndexTestModel )
187+ index_info = connection .introspection .get_constraints (
188+ cursor = None ,
189+ table_name = SearchIndexTestModel ._meta .db_table ,
190+ )
191+ expected_options = {
192+ "latestDefinition" : {
193+ "fields" : [
194+ {"path" : "text" , "type" : "filter" },
195+ {"path" : "object_id" , "type" : "filter" },
196+ {"path" : "number" , "type" : "filter" },
197+ {"path" : "embedded" , "type" : "filter" },
198+ {
199+ "numDimensions" : 10 ,
200+ "path" : "vector_integer" ,
201+ "similarity" : "cosine" ,
202+ "type" : "vector" ,
203+ },
204+ {
205+ "numDimensions" : 10 ,
206+ "path" : "vector_float" ,
207+ "similarity" : "cosine" ,
208+ "type" : "vector" ,
209+ },
210+ {"path" : "boolean" , "type" : "filter" },
211+ {"path" : "date" , "type" : "filter" },
212+ ]
213+ },
214+ "latestVersion" : 0 ,
215+ "name" : "recent_test_idx" ,
216+ "queryable" : False ,
217+ "type" : "vectorSearch" ,
218+ }
219+ self .assertCountEqual (index_info [index .name ]["columns" ], index .fields )
220+ index_info [index .name ]["options" ].pop ("id" )
221+ index_info [index .name ]["options" ].pop ("status" )
222+ self .assertEqual (index_info [index .name ]["options" ], expected_options )
223+ self .assertExistAndRemoveIndex (index = index , model = SearchIndexTestModel )
0 commit comments