Improved tests [skip ci]

Author: ankane

tests/test_django.py

@@ -165,7 +165,7 @@ def setup_method(self):
     def test_vector(self):
         Item(id=1, embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)
-        assert np.array_equal(item.embedding, np.array([1, 2, 3]))
+        assert np.array_equal(item.embedding, [1, 2, 3])
         assert item.embedding.dtype == np.float32
 
     def test_vector_l2_distance(self):
@@ -293,15 +293,15 @@ def test_vector_avg(self):
         Item(embedding=[1, 2, 3]).save()
         Item(embedding=[4, 5, 6]).save()
         avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg']
-        assert np.array_equal(avg, np.array([2.5, 3.5, 4.5]))
+        assert np.array_equal(avg, [2.5, 3.5, 4.5])
 
     def test_vector_sum(self):
         sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum']
         assert sum is None
         Item(embedding=[1, 2, 3]).save()
         Item(embedding=[4, 5, 6]).save()
         sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum']
-        assert np.array_equal(sum, np.array([5, 7, 9]))
+        assert np.array_equal(sum, [5, 7, 9])
 
     def test_halfvec_avg(self):
         avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg']
@@ -347,7 +347,7 @@ def test_vector_form_save(self):
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist()
+        assert np.array_equal(Item.objects.get(pk=1).embedding, [4, 5, 6])
 
     def test_vector_form_save_missing(self):
         Item(id=1).save()
@@ -465,8 +465,8 @@ def test_vector_array(self):
 
             # this fails if the driver does not cast arrays
             item = Item.objects.get(pk=1)
-            assert item.embeddings[0].tolist() == [1, 2, 3]
-            assert item.embeddings[1].tolist() == [4, 5, 6]
+            assert np.array_equal(item.embeddings[0], [1, 2, 3])
+            assert np.array_equal(item.embeddings[1], [4, 5, 6])
 
     def test_double_array(self):
         Item(id=1, double_embedding=[1, 1, 1]).save()

tests/test_half_vector.py

@@ -49,11 +49,11 @@ def test_dimensions(self):
     def test_from_text(self):
         vec = HalfVector.from_text('[1.5,2,3]')
         assert vec.to_list() == [1.5, 2, 3]
-        assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3]))
+        assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
 
     def test_from_binary(self):
         data = pack('>HH3e', 3, 0, 1.5, 2, 3)
         vec = HalfVector.from_binary(data)
         assert vec.to_list() == [1.5, 2, 3]
-        assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3]))
+        assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
         assert vec.to_binary() == data

tests/test_peewee.py

@@ -43,7 +43,7 @@ def setup_method(self):
     def test_vector(self):
         Item.create(id=1, embedding=[1, 2, 3])
         item = Item.get_by_id(1)
-        assert np.array_equal(item.embedding, np.array([1, 2, 3]))
+        assert np.array_equal(item.embedding, [1, 2, 3])
         assert item.embedding.dtype == np.float32
 
     def test_vector_l2_distance(self):
@@ -170,15 +170,15 @@ def test_vector_avg(self):
         Item.create(embedding=[1, 2, 3])
         Item.create(embedding=[4, 5, 6])
         avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar()
-        assert np.array_equal(avg, np.array([2.5, 3.5, 4.5]))
+        assert np.array_equal(avg, [2.5, 3.5, 4.5])
 
     def test_vector_sum(self):
         sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar()
         assert sum is None
         Item.create(embedding=[1, 2, 3])
         Item.create(embedding=[4, 5, 6])
         sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar()
-        assert np.array_equal(sum, np.array([5, 7, 9]))
+        assert np.array_equal(sum, [5, 7, 9])
 
     def test_halfvec_avg(self):
         avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar()
@@ -220,5 +220,5 @@ class Meta:
         # fails with column "embeddings" is of type vector[] but expression is of type text[]
         # ExtItem.create(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])
         # item = ExtItem.get_by_id(1)
-        # assert np.array_equal(item.embeddings[0], np.array([1, 2, 3]))
-        # assert np.array_equal(item.embeddings[1], np.array([4, 5, 6]))
+        # assert np.array_equal(item.embeddings[0], [1, 2, 3])
+        # assert np.array_equal(item.embeddings[1], [4, 5, 6])

tests/test_psycopg.py

@@ -46,23 +46,23 @@ def test_vector_text_format_non_contiguous(self):
         embedding = np.flipud(np.array([1.5, 2, 3]))
         assert not embedding.data.contiguous
         res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0]
-        assert np.array_equal(res, np.array([3, 2, 1.5]))
+        assert np.array_equal(res, [3, 2, 1.5])
 
     def test_vector_binary_format_non_contiguous(self):
         embedding = np.flipud(np.array([1.5, 2, 3]))
         assert not embedding.data.contiguous
         res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0]
-        assert np.array_equal(res, np.array([3, 2, 1.5]))
+        assert np.array_equal(res, [3, 2, 1.5])
 
     def test_vector_class_binary_format(self):
         embedding = Vector([1.5, 2, 3])
         res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0]
-        assert np.array_equal(res, np.array([1.5, 2, 3]))
+        assert np.array_equal(res, [1.5, 2, 3])
 
     def test_vector_class_text_format(self):
         embedding = Vector([1.5, 2, 3])
         res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0]
-        assert np.array_equal(res, np.array([1.5, 2, 3]))
+        assert np.array_equal(res, [1.5, 2, 3])
 
     def test_halfvec(self):
         embedding = HalfVector([1.5, 2, 3])
@@ -182,7 +182,7 @@ def configure(conn):
 
         with pool.connection() as conn:
             res = conn.execute("SELECT '[1,2,3]'::vector").fetchone()
-            assert np.array_equal(res[0], np.array([1, 2, 3]))
+            assert np.array_equal(res[0], [1, 2, 3])
 
         pool.close()
 
@@ -218,6 +218,6 @@ async def configure(conn):
             async with conn.cursor() as cur:
                 await cur.execute("SELECT '[1,2,3]'::vector")
                 res = await cur.fetchone()
-                assert np.array_equal(res[0], np.array([1, 2, 3]))
+                assert np.array_equal(res[0], [1, 2, 3])
 
         await pool.close()

tests/test_psycopg2.py

@@ -122,7 +122,7 @@ def test_pool(self):
             cur = conn.cursor()
             cur.execute("SELECT '[1,2,3]'::vector")
             res = cur.fetchone()
-            assert np.array_equal(res[0], np.array([1, 2, 3]))
+            assert np.array_equal(res[0], [1, 2, 3])
         finally:
             pool.putconn(conn)
 

tests/test_sparse_vector.py

@@ -9,7 +9,7 @@ class TestSparseVector:
     def test_list(self):
         vec = SparseVector([1, 0, 2, 0, 3, 0])
         assert vec.to_list() == [1, 0, 2, 0, 3, 0]
-        assert vec.to_numpy().tolist() == [1, 0, 2, 0, 3, 0]
+        assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0])
         assert vec.indices() == [0, 2, 4]
 
     def test_list_dimensions(self):
@@ -69,7 +69,7 @@ def test_values(self):
         assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3]
 
     def test_to_coo(self):
-        assert SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]]
+        assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]])
 
     def test_zero_vector_text(self):
         vec = SparseVector({}, 3)
@@ -81,7 +81,7 @@ def test_from_text(self):
         assert vec.indices() == [0, 2, 4]
         assert vec.values() == [1.5, 2, 3]
         assert vec.to_list() == [1.5, 0, 2, 0, 3, 0]
-        assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0]))
+        assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0])
 
     def test_from_binary(self):
         data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3)
@@ -90,5 +90,5 @@ def test_from_binary(self):
         assert vec.indices() == [0, 2, 4]
         assert vec.values() == [1.5, 2, 3]
         assert vec.to_list() == [1.5, 0, 2, 0, 3, 0]
-        assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0]))
+        assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0])
         assert vec.to_binary() == data

tests/test_sqlalchemy.py

@@ -201,7 +201,7 @@ def test_vector(self, engine):
             session.add(Item(id=1, embedding=[1, 2, 3]))
             session.commit()
             item = session.get(Item, 1)
-            assert item.embedding.tolist() == [1, 2, 3]
+            assert np.array_equal(item.embedding, [1, 2, 3])
 
     def test_vector_l2_distance(self, engine):
         create_items()
@@ -509,7 +509,7 @@ def test_automap(self, engine):
         with Session(engine) as session:
             session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}])
             item = session.query(AutoItem).first()
-            assert item.embedding.tolist() == [1, 2, 3]
+            assert np.array_equal(item.embedding, [1, 2, 3])
 
     def test_half_precision(self, engine):
         create_items()
@@ -541,8 +541,8 @@ def test_vector_array(self, engine):
 
             # this fails if the driver does not cast arrays
             item = session.get(Item, 1)
-            assert item.embeddings[0].tolist() == [1, 2, 3]
-            assert item.embeddings[1].tolist() == [4, 5, 6]
+            assert np.array_equal(item.embeddings[0], [1, 2, 3])
+            assert np.array_equal(item.embeddings[1], [4, 5, 6])
 
     def test_halfvec_array(self, engine):
         with Session(engine) as session:
@@ -621,7 +621,7 @@ async def test_avg(self, engine):
                 session.add(Item(embedding=[1, 2, 3]))
                 session.add(Item(embedding=[4, 5, 6]))
                 res = await session.scalars(select(avg(Item.embedding)))
-                assert res.first().tolist() == [2.5, 3.5, 4.5]
+                assert np.array_equal(res.first(), [2.5, 3.5, 4.5])
 
         await engine.dispose()
 
@@ -639,12 +639,12 @@ async def test_vector_array(self, engine):
             async with session.begin():
                 session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]))
                 item = await session.get(Item, 1)
-                assert item.embeddings[0].tolist() == [1, 2, 3]
-                assert item.embeddings[1].tolist() == [4, 5, 6]
+                assert np.array_equal(item.embeddings[0], [1, 2, 3])
+                assert np.array_equal(item.embeddings[1], [4, 5, 6])
 
                 session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
                 item = await session.get(Item, 2)
-                assert item.embeddings[0].tolist() == [1, 2, 3]
-                assert item.embeddings[1].tolist() == [4, 5, 6]
+                assert np.array_equal(item.embeddings[0], [1, 2, 3])
+                assert np.array_equal(item.embeddings[1], [4, 5, 6])
 
         await engine.dispose()

tests/test_sqlmodel.py

@@ -76,7 +76,7 @@ def test_vector(self):
             session.add(Item(id=1, embedding=[1, 2, 3]))
             session.commit()
             item = session.get(Item, 1)
-            assert item.embedding.tolist() == [1, 2, 3]
+            assert np.array_equal(item.embedding, np.array([1, 2, 3]))
 
     def test_vector_l2_distance(self):
         create_items()

tests/test_vector.py

@@ -49,11 +49,11 @@ def test_dimensions(self):
     def test_from_text(self):
         vec = Vector.from_text('[1.5,2,3]')
         assert vec.to_list() == [1.5, 2, 3]
-        assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3]))
+        assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
 
     def test_from_binary(self):
         data = pack('>HH3f', 3, 0, 1.5, 2, 3)
         vec = Vector.from_binary(data)
         assert vec.to_list() == [1.5, 2, 3]
-        assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3]))
+        assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
         assert vec.to_binary() == data