Improved asyncpg tests [skip ci]

Author: ankane

tests/test_asyncpg.py

@@ -1,6 +1,6 @@
 import asyncpg
 import numpy as np
-from pgvector import SparseVector
+from pgvector import HalfVector, SparseVector, Vector
 from pgvector.asyncpg import register_vector
 import pytest
 
@@ -15,13 +15,15 @@ async def test_vector(self):
 
         await register_vector(conn)
 
-        embedding = np.array([1.5, 2, 3])
-        await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
+        embedding = Vector([1.5, 2, 3])
+        embedding2 = np.array([4.5, 5, 6])
+        await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert np.array_equal(res[0]['embedding'], embedding)
+        assert np.array_equal(res[0]['embedding'], embedding.to_numpy())
         assert res[0]['embedding'].dtype == np.float32
-        assert res[1]['embedding'] is None
+        assert np.array_equal(res[1]['embedding'], embedding2)
+        assert res[2]['embedding'] is None
 
         # ensures binary format is correct
         text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1")
@@ -38,12 +40,14 @@ async def test_halfvec(self):
 
         await register_vector(conn)
 
-        embedding = [1.5, 2, 3]
-        await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
+        embedding = HalfVector([1.5, 2, 3])
+        embedding2 = [4.5, 5, 6]
+        await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert res[0]['embedding'].to_list() == [1.5, 2, 3]
-        assert res[1]['embedding'] is None
+        assert res[0]['embedding'] == embedding
+        assert res[1]['embedding'] == HalfVector(embedding2)
+        assert res[2]['embedding'] is None
 
         # ensures binary format is correct
         text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1")
@@ -87,7 +91,7 @@ async def test_sparsevec(self):
         await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert res[0]['embedding'].to_list() == [1.5, 2, 3]
+        assert res[0]['embedding'] == embedding
         assert res[1]['embedding'] is None
 
         # ensures binary format is correct
@@ -105,12 +109,15 @@ async def test_vector_array(self):
 
         await register_vector(conn)
 
-        embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
-        await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings[0], embeddings[1])
+        embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])]
+        embeddings2 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
+        await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[]), (ARRAY[$3, $4]::vector[])", embeddings[0], embeddings[1], embeddings2[0], embeddings2[1])
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert np.array_equal(res[0]['embeddings'][0], embeddings[0])
-        assert np.array_equal(res[0]['embeddings'][1], embeddings[1])
+        assert np.array_equal(res[0]['embeddings'][0], embeddings[0].to_numpy())
+        assert np.array_equal(res[0]['embeddings'][1], embeddings[1].to_numpy())
+        assert np.array_equal(res[1]['embeddings'][0], embeddings2[0])
+        assert np.array_equal(res[1]['embeddings'][1], embeddings2[1])
 
         await conn.close()
 
@@ -126,10 +133,12 @@ async def init(conn):
             await conn.execute('DROP TABLE IF EXISTS asyncpg_items')
             await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))')
 
-            embedding = np.array([1.5, 2, 3])
-            await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
+            embedding = Vector([1.5, 2, 3])
+            embedding2 = np.array([1.5, 2, 3])
+            await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2)
 
             res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-            assert np.array_equal(res[0]['embedding'], embedding)
+            assert np.array_equal(res[0]['embedding'], embedding.to_numpy())
             assert res[0]['embedding'].dtype == np.float32
-            assert res[1]['embedding'] is None
+            assert np.array_equal(res[1]['embedding'], embedding2)
+            assert res[2]['embedding'] is None