Improved array tests

Author: ankane

tests/test_sqlalchemy.py

@@ -27,10 +27,11 @@
     session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
     session.commit()
 
-array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')
+psycopg2_array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')
+array_engines = [psycopg2_array_engine]
 
 
-@event.listens_for(array_engine, "connect")
+@event.listens_for(psycopg2_array_engine, "connect")
 def connect(dbapi_connection, connection_record):
     from pgvector.psycopg2 import register_vector
     register_vector(dbapi_connection, globally=False, arrays=True)
@@ -468,8 +469,31 @@ def test_automap(self, engine):
             item = session.query(AutoItem).first()
             assert item.embedding.tolist() == [1, 2, 3]
 
+    def test_half_precision(self, engine):
+        create_items()
+        with Session(engine) as session:
+            items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all()
+            assert [v.id for v in items] == [1, 3, 2]
+
+    def test_binary_quantize(self, engine):
+        with Session(engine) as session:
+            session.add(Item(id=1, embedding=[-1, -2, -3]))
+            session.add(Item(id=2, embedding=[1, -2, 3]))
+            session.add(Item(id=3, embedding=[1, 2, 3]))
+            session.commit()
+
+            distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+            items = session.query(Item).order_by(distance).all()
+            assert [v.id for v in items] == [2, 3, 1]
+
+
+@pytest.mark.parametrize("engine", array_engines)
+class TestSqlalchemyArray:
+    def setup_method(self):
+        delete_items()
+
     def test_vector_array(self, engine):
-        with Session(array_engine) as session:
+        with Session(engine) as session:
             session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
             session.commit()
 
@@ -479,7 +503,7 @@ def test_vector_array(self, engine):
             assert item.embeddings[1].tolist() == [4, 5, 6]
 
     def test_halfvec_array(self, engine):
-        with Session(array_engine) as session:
+        with Session(engine) as session:
             session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
             session.commit()
 
@@ -488,23 +512,6 @@ def test_halfvec_array(self, engine):
             assert item.half_embeddings[0].to_list() == [1, 2, 3]
             assert item.half_embeddings[1].to_list() == [4, 5, 6]
 
-    def test_half_precision(self, engine):
-        create_items()
-        with Session(engine) as session:
-            items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all()
-            assert [v.id for v in items] == [1, 3, 2]
-
-    def test_binary_quantize(self, engine):
-        with Session(engine) as session:
-            session.add(Item(id=1, embedding=[-1, -2, -3]))
-            session.add(Item(id=2, embedding=[1, -2, 3]))
-            session.add(Item(id=3, embedding=[1, 2, 3]))
-            session.commit()
-
-            distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
-            items = session.query(Item).order_by(distance).all()
-            assert [v.id for v in items] == [2, 3, 1]
-
 
 class TestSqlalchemyAsync:
     def setup_method(self):