Small refactor of unit tests to consolidate language and add a query …

…and check helper to `test_ingestion.py` (#176)
TileDB-Inc · Dec 21, 2023 · 9c029e9 · 9c029e9
1 parent f5734f9
commit 9c029e9
Show file tree

Hide file tree

Showing 3 changed files with 131 additions and 138 deletions.
diff --git a/apis/python/test/test_cloud.py b/apis/python/test/test_cloud.py
@@ -40,7 +40,7 @@ def test_cloud_flat(self):
         k = 100
         nqueries = 100
 
-        query_vectors = load_fvecs(queries_uri)
+        queries = load_fvecs(queries_uri)
         gt_i, gt_d = get_groundtruth_ivec(gt_uri, k=k, nqueries=nqueries)
 
         index = vs.ingest(
@@ -50,7 +50,7 @@ def test_cloud_flat(self):
             config=tiledb.cloud.Config().dict(),
             mode=Mode.BATCH,
         )
-        _, result_i = index.query(query_vectors, k=k)
+        _, result_i = index.query(queries, k=k)
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
 
     def test_cloud_ivf_flat(self):
@@ -63,7 +63,7 @@ def test_cloud_ivf_flat(self):
         nqueries = 100
         nprobe = 20
 
-        query_vectors = load_fvecs(queries_uri)
+        queries = load_fvecs(queries_uri)
         gt_i, gt_d = get_groundtruth_ivec(gt_uri, k=k, nqueries=nqueries)
 
         index = vs.ingest(
@@ -79,16 +79,16 @@ def test_cloud_ivf_flat(self):
             # mode=Mode.BATCH,
         )
 
-        _, result_i = index.query(query_vectors, k=k, nprobe=nprobe)
+        _, result_i = index.query(queries, k=k, nprobe=nprobe)
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
 
         _, result_i = index.query(
-            query_vectors, k=k, nprobe=nprobe, mode=Mode.REALTIME, num_partitions=2, resource_class="standard"
+            queries, k=k, nprobe=nprobe, mode=Mode.REALTIME, num_partitions=2, resource_class="standard"
         )
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
 
         _, result_i = index.query(
-            query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL, num_partitions=2
+            queries, k=k, nprobe=nprobe, mode=Mode.LOCAL, num_partitions=2
         )
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
 
@@ -97,20 +97,20 @@ def test_cloud_ivf_flat(self):
 
         # Cannot pass resource_class or resources to LOCAL mode or to no mode.
         with self.assertRaises(TypeError):
-            index.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL, resource_class="large")
+            index.query(queries, k=k, nprobe=nprobe, mode=Mode.LOCAL, resource_class="large")
         with self.assertRaises(TypeError):
-            index.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL, resources=resources)
+            index.query(queries, k=k, nprobe=nprobe, mode=Mode.LOCAL, resources=resources)
         with self.assertRaises(TypeError):
-            index.query(query_vectors, k=k, nprobe=nprobe, resource_class="large")
+            index.query(queries, k=k, nprobe=nprobe, resource_class="large")
         with self.assertRaises(TypeError):
-            index.query(query_vectors, k=k, nprobe=nprobe, resources=resources)
+            index.query(queries, k=k, nprobe=nprobe, resources=resources)
 
         # Cannot pass resources to REALTIME.
         with self.assertRaises(TypeError):
-            index.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.REALTIME, resources=resources)
+            index.query(queries, k=k, nprobe=nprobe, mode=Mode.REALTIME, resources=resources)
 
         # Cannot pass both resource_class and resources.
         with self.assertRaises(TypeError):
-            index.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.REALTIME, resource_class="large", resources=resources)
+            index.query(queries, k=k, nprobe=nprobe, mode=Mode.REALTIME, resource_class="large", resources=resources)
         with self.assertRaises(TypeError):
-            index.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.BATCH, resource_class="large", resources=resources)
+            index.query(queries, k=k, nprobe=nprobe, mode=Mode.BATCH, resource_class="large", resources=resources)
diff --git a/apis/python/test/test_index.py b/apis/python/test/test_index.py
@@ -131,8 +131,8 @@ def test_index_with_incorrect_num_of_query_columns_simple(tmp_path):
             index.query(np.random.rand(*query_shape).astype(np.float32), k=10)
 
         # Okay otherwise.
-        query_vectors = load_fvecs(queries_uri)
-        index.query(query_vectors, k=10)
+        queries = load_fvecs(queries_uri)
+        index.query(queries, k=10)
 
 def test_index_with_incorrect_num_of_query_columns_complex(tmp_path):
     # Tests that we raise a TypeError if the number of columns in the query is not the same as the