Skip to content

Commit

Permalink
Added test for arrays with asyncpg
Browse files Browse the repository at this point in the history
  • Loading branch information
ankane committed Oct 13, 2024
1 parent 82b7ab9 commit 7887a34
Showing 1 changed file with 18 additions and 0 deletions.
18 changes: 18 additions & 0 deletions tests/test_asyncpg.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,24 @@ async def test_sparsevec(self):

await conn.close()

@pytest.mark.asyncio
async def test_vector_array(self):
conn = await asyncpg.connect(database='pgvector_python_test')
await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
await conn.execute('DROP TABLE IF EXISTS asyncpg_items')
await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])')

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])

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])

await conn.close()

@pytest.mark.asyncio
async def test_pool(self):
async def init(conn):
Expand Down

0 comments on commit 7887a34

Please sign in to comment.