Skip to content

Commit

Permalink
Added more tests for SQLAlchemy
Browse files Browse the repository at this point in the history
  • Loading branch information
ankane committed Feb 9, 2025
1 parent 00cd08e commit 7837e92
Showing 1 changed file with 30 additions and 0 deletions.
30 changes: 30 additions & 0 deletions tests/test_sqlalchemy.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ def connect(dbapi_connection, connection_record):
register_vector(dbapi_connection, globally=False, arrays=True)


psycopg3_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test')

Base = declarative_base()


Expand Down Expand Up @@ -493,6 +495,34 @@ def test_binary_quantize(self):
items = session.query(Item).order_by(distance).all()
assert [v.id for v in items] == [2, 3, 1]

def test_psycopg_vector(self):
with Session(psycopg3_engine) as session:
session.add(Item(id=1, embedding=[1, 2, 3]))
session.commit()
item = session.get(Item, 1)
assert item.embedding.tolist() == [1, 2, 3]

def test_psycopg_halfvec(self):
with Session(psycopg3_engine) as session:
session.add(Item(id=1, half_embedding=[1, 2, 3]))
session.commit()
item = session.get(Item, 1)
assert item.half_embedding.to_list() == [1, 2, 3]

def test_psycopg_bit(self):
with Session(psycopg3_engine) as session:
session.add(Item(id=1, binary_embedding='101'))
session.commit()
item = session.get(Item, 1)
assert item.binary_embedding == '101'

def test_psycopg_sparsevec(self):
with Session(psycopg3_engine) as session:
session.add(Item(id=1, sparse_embedding=[1, 2, 3]))
session.commit()
item = session.get(Item, 1)
assert item.sparse_embedding.to_list() == [1, 2, 3]

@pytest.mark.asyncio
@pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
async def test_psycopg_async_avg(self):
Expand Down

0 comments on commit 7837e92

Please sign in to comment.