Improved tests [skip ci]

Author: ankane

tests/test_half_vector.py

@@ -1,6 +1,7 @@
 import numpy as np
 from pgvector import HalfVector
 import pytest
+from struct import pack
 
 
 class TestHalfVector:
@@ -49,3 +50,10 @@ 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]))
+
+    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 vec.to_binary() == data

tests/test_psycopg.py

@@ -75,9 +75,6 @@ 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])
-        # TODO move
-        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])
@@ -113,12 +110,6 @@ 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
-        # TODO move
-        assert res.dimensions() == 6
-        assert res.indices() == [0, 2, 4]
-        assert res.values() == [1.5, 2, 3]
-        assert res.to_list() == [1.5, 0, 2, 0, 3, 0]
-        assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0]))
 
     def test_sparsevec_text_format(self):
         embedding = SparseVector([1.5, 0, 2, 0, 3, 0])

tests/test_sparse_vector.py

@@ -2,6 +2,7 @@
 from pgvector import SparseVector
 import pytest
 from scipy.sparse import coo_array
+from struct import pack
 
 
 class TestSparseVector:
@@ -81,3 +82,13 @@ def test_from_text(self):
         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]))
+
+    def test_from_binary(self):
+        data = pack('>iii3i3f', 6, 3, 0, *[0, 2, 4], *[1.5, 2, 3])
+        vec = SparseVector.from_binary(data)
+        assert vec.dimensions() == 6
+        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 vec.to_binary() == data

tests/test_vector.py

@@ -1,6 +1,7 @@
 import numpy as np
 from pgvector import Vector
 import pytest
+from struct import pack
 
 
 class TestVector:
@@ -44,3 +45,15 @@ def test_equality(self):
 
     def test_dimensions(self):
         assert Vector([1, 2, 3]).dimensions() == 3
+
+    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]))
+
+    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 vec.to_binary() == data