Skip to content

Commit 7837e92

Browse files
committed
Added more tests for SQLAlchemy
1 parent 00cd08e commit 7837e92

File tree

1 file changed

+30
-0
lines changed

1 file changed

+30
-0
lines changed

Diff for: tests/test_sqlalchemy.py

+30
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ def connect(dbapi_connection, connection_record):
2929
register_vector(dbapi_connection, globally=False, arrays=True)
3030

3131

32+
psycopg3_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test')
33+
3234
Base = declarative_base()
3335

3436

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

498+
def test_psycopg_vector(self):
499+
with Session(psycopg3_engine) as session:
500+
session.add(Item(id=1, embedding=[1, 2, 3]))
501+
session.commit()
502+
item = session.get(Item, 1)
503+
assert item.embedding.tolist() == [1, 2, 3]
504+
505+
def test_psycopg_halfvec(self):
506+
with Session(psycopg3_engine) as session:
507+
session.add(Item(id=1, half_embedding=[1, 2, 3]))
508+
session.commit()
509+
item = session.get(Item, 1)
510+
assert item.half_embedding.to_list() == [1, 2, 3]
511+
512+
def test_psycopg_bit(self):
513+
with Session(psycopg3_engine) as session:
514+
session.add(Item(id=1, binary_embedding='101'))
515+
session.commit()
516+
item = session.get(Item, 1)
517+
assert item.binary_embedding == '101'
518+
519+
def test_psycopg_sparsevec(self):
520+
with Session(psycopg3_engine) as session:
521+
session.add(Item(id=1, sparse_embedding=[1, 2, 3]))
522+
session.commit()
523+
item = session.get(Item, 1)
524+
assert item.sparse_embedding.to_list() == [1, 2, 3]
525+
496526
@pytest.mark.asyncio
497527
@pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
498528
async def test_psycopg_async_avg(self):

0 commit comments

Comments
 (0)