Skip to content

Commit 0d6f719

Browse files
timgrahamtimgraham
committed
reuse assertAddRemoveIndex
1 parent 3f43d46 commit 0d6f719

File tree

3 files changed

+101
-105
lines changed

3 files changed

+101
-105
lines changed

‎tests/indexes_/test_base.py

+23
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
from django.db import connection
2+
3+
4+
class SchemaAssertionMixin:
5+
def assertAddRemoveIndex(self, editor, model, index):
6+
editor.add_index(index=index, model=model)
7+
try:
8+
self.assertIn(
9+
index.name,
10+
connection.introspection.get_constraints(
11+
cursor=None,
12+
table_name=model._meta.db_table,
13+
),
14+
)
15+
finally:
16+
editor.remove_index(index=index, model=model)
17+
self.assertNotIn(
18+
index.name,
19+
connection.introspection.get_constraints(
20+
cursor=None,
21+
table_name=model._meta.db_table,
22+
),
23+
)

‎tests/indexes_/test_condition.py

+2-12
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,10 @@
55
from django.test import TestCase
66

77
from .models import Article
8+
from .test_base import SchemaAssertionMixin
89

910

10-
class PartialIndexTests(TestCase):
11-
def assertAddRemoveIndex(self, editor, model, index):
12-
editor.add_index(index=index, model=model)
13-
self.assertIn(
14-
index.name,
15-
connection.introspection.get_constraints(
16-
cursor=None,
17-
table_name=model._meta.db_table,
18-
),
19-
)
20-
editor.remove_index(index=index, model=model)
21-
11+
class PartialIndexTests(SchemaAssertionMixin, TestCase):
2212
def test_not_supported(self):
2313
msg = "MongoDB does not support the 'isnull' lookup in indexes."
2414
with connection.schema_editor() as editor, self.assertRaisesMessage(NotSupportedError, msg):

‎tests/indexes_/test_search_indexes.py

+76-93
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from django_mongodb_backend.indexes import SearchIndex, VectorSearchIndex
55

66
from .models import SearchIndexTestModel
7+
from .test_base import SchemaAssertionMixin
78

89

910
@skipIfDBFeature("supports_atlas_search")
@@ -79,36 +80,14 @@ def test_define_field_twice(self):
7980
)
8081

8182

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-
10283
@skipUnlessDBFeature("supports_atlas_search")
10384
class SearchIndexSchemaTests(SchemaAssertionMixin, TestCase):
10485
def test_simple(self):
10586
index = SearchIndex(
10687
name="recent_test_idx",
10788
fields=["text"],
10889
)
109-
with connection.schema_editor() as editor:
110-
editor.add_index(index=index, model=SearchIndexTestModel)
111-
self.assertExistAndRemoveIndex(index=index, model=SearchIndexTestModel)
90+
self.assertAddRemoveIndex(index=index, model=SearchIndexTestModel)
11291

11392
def test_all_fields(self):
11493
index = SearchIndex(
@@ -127,46 +106,47 @@ def test_all_fields(self):
127106
)
128107
with connection.schema_editor() as editor:
129108
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",
109+
try:
110+
index_info = connection.introspection.get_constraints(
111+
cursor=None,
112+
table_name=SearchIndexTestModel._meta.db_table,
113+
)
114+
expected_options = {
115+
"dynamic": False,
116+
"fields": {
117+
"boolean": {"type": "boolean"},
118+
"embedded": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
119+
"json_data": {"dynamic": False, "fields": {}, "type": "document"},
120+
"number": {
121+
"indexDoubles": True,
122+
"indexIntegers": True,
123+
"representation": "double",
124+
"type": "number",
125+
},
126+
"object_id": {"type": "objectId"},
127+
"text": {
128+
"indexOptions": "offsets",
129+
"norms": "include",
130+
"store": True,
131+
"type": "string",
132+
},
133+
"date": {"type": "date"},
134+
"vector_float": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
135+
"vector_integer": {"dynamic": False, "fields": {}, "type": "embeddedDocuments"},
152136
},
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)
137+
}
138+
self.assertCountEqual(index_info[index.name]["columns"], index.fields)
139+
self.assertEqual(index_info[index.name]["options"], expected_options)
140+
finally:
141+
with connection.schema_editor() as editor:
142+
editor.remove_index(index=index, model=SearchIndexTestModel)
161143

162144

163145
@skipUnlessDBFeature("supports_atlas_search")
164146
class VectorSearchIndexSchemaTests(SchemaAssertionMixin, TestCase):
165147
def test_simple_vector_search(self):
166148
index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
167-
with connection.schema_editor() as editor:
168-
editor.add_index(index=index, model=SearchIndexTestModel)
169-
self.assertExistAndRemoveIndex(index=index, model=SearchIndexTestModel)
149+
self.assertAddRemoveIndex(index=index, model=SearchIndexTestModel)
170150

171151
def test_multiple_fields(self):
172152
index = VectorSearchIndex(
@@ -184,40 +164,43 @@ def test_multiple_fields(self):
184164
)
185165
with connection.schema_editor() as editor:
186166
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)
167+
try:
168+
index_info = connection.introspection.get_constraints(
169+
cursor=None,
170+
table_name=SearchIndexTestModel._meta.db_table,
171+
)
172+
expected_options = {
173+
"latestDefinition": {
174+
"fields": [
175+
{"path": "text", "type": "filter"},
176+
{"path": "object_id", "type": "filter"},
177+
{"path": "number", "type": "filter"},
178+
{"path": "embedded", "type": "filter"},
179+
{
180+
"numDimensions": 10,
181+
"path": "vector_integer",
182+
"similarity": "cosine",
183+
"type": "vector",
184+
},
185+
{
186+
"numDimensions": 10,
187+
"path": "vector_float",
188+
"similarity": "cosine",
189+
"type": "vector",
190+
},
191+
{"path": "boolean", "type": "filter"},
192+
{"path": "date", "type": "filter"},
193+
]
194+
},
195+
"latestVersion": 0,
196+
"name": "recent_test_idx",
197+
"queryable": False,
198+
"type": "vectorSearch",
199+
}
200+
self.assertCountEqual(index_info[index.name]["columns"], index.fields)
201+
index_info[index.name]["options"].pop("id")
202+
index_info[index.name]["options"].pop("status")
203+
self.assertEqual(index_info[index.name]["options"], expected_options)
204+
finally:
205+
with connection.schema_editor() as editor:
206+
editor.remove_index(index=index, model=SearchIndexTestModel)

0 commit comments

Comments
 (0)