Improved tests [skip ci]

Author: ankane

tests/test_django.py

@@ -199,7 +199,7 @@ def test_vector_l1_distance(self):
     def test_halfvec(self):
         Item(id=1, half_embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)
-        assert item.half_embedding.to_list() == [1, 2, 3]
+        assert item.half_embedding == HalfVector([1, 2, 3])
 
     def test_halfvec_l2_distance(self):
         create_items()
@@ -251,7 +251,7 @@ def test_bit_jaccard_distance(self):
     def test_sparsevec(self):
         Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save()
         item = Item.objects.get(pk=1)
-        assert item.sparse_embedding.to_list() == [1, 2, 3]
+        assert item.sparse_embedding == SparseVector([1, 2, 3])
 
     def test_sparsevec_l2_distance(self):
         create_items()
@@ -309,15 +309,15 @@ def test_halfvec_avg(self):
         Item(half_embedding=[1, 2, 3]).save()
         Item(half_embedding=[4, 5, 6]).save()
         avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg']
-        assert avg.to_list() == [2.5, 3.5, 4.5]
+        assert avg == HalfVector([2.5, 3.5, 4.5])
 
     def test_halfvec_sum(self):
         sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum']
         assert sum is None
         Item(half_embedding=[1, 2, 3]).save()
         Item(half_embedding=[4, 5, 6]).save()
         sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum']
-        assert sum.to_list() == [5, 7, 9]
+        assert sum == HalfVector([5, 7, 9])
 
     def test_serialization(self):
         create_items()
@@ -375,7 +375,7 @@ def test_halfvec_form_save(self):
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert [4, 5, 6] == Item.objects.get(pk=1).half_embedding.to_list()
+        assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6])
 
     def test_halfvec_form_save_missing(self):
         Item(id=1).save()
@@ -432,7 +432,7 @@ def test_sparsevec_form_save(self):
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert [4, 5, 6] == Item.objects.get(pk=1).sparse_embedding.to_list()
+        assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6])
 
     def test_sparesevec_form_save_missing(self):
         Item(id=1).save()

tests/test_psycopg.py

@@ -69,17 +69,19 @@ def test_halfvec(self):
         conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,))
 
         res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0]
-        assert res.to_list() == [1.5, 2, 3]
+        assert res == HalfVector([1.5, 2, 3])
 
     def test_halfvec_binary_format(self):
         embedding = HalfVector([1.5, 2, 3])
         res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0]
+        assert res == HalfVector([1.5, 2, 3])
         assert res.to_list() == [1.5, 2, 3]
         assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3]))
 
     def test_halfvec_text_format(self):
         embedding = HalfVector([1.5, 2, 3])
         res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0]
+        assert res == HalfVector([1.5, 2, 3])
         assert res.to_list() == [1.5, 2, 3]
         assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3]))
 
@@ -106,11 +108,12 @@ def test_sparsevec(self):
         conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,))
 
         res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0]
-        assert res.to_list() == [1.5, 2, 3]
+        assert res == SparseVector([1.5, 2, 3])
 
     def test_sparsevec_binary_format(self):
         embedding = SparseVector([1.5, 0, 2, 0, 3, 0])
         res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0]
+        assert res == embedding
         assert res.dimensions() == 6
         assert res.indices() == [0, 2, 4]
         assert res.values() == [1.5, 2, 3]
@@ -120,6 +123,7 @@ def test_sparsevec_binary_format(self):
     def test_sparsevec_text_format(self):
         embedding = SparseVector([1.5, 0, 2, 0, 3, 0])
         res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0]
+        assert res == embedding
         assert res.dimensions() == 6
         assert res.indices() == [0, 2, 4]
         assert res.values() == [1.5, 2, 3]