From 70ae49a5e2e7e6c393f3b37041049091cf3007d5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 10:48:09 -0700 Subject: [PATCH] Updated tests to use globally=False --- tests/test_peewee.py | 12 ++++++++---- tests/test_psycopg2.py | 6 +++--- tests/test_sqlalchemy.py | 6 ++++-- tests/test_sqlmodel.py | 6 ++++-- 4 files changed, 19 insertions(+), 11 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 0882890..72502db 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -169,7 +169,8 @@ def test_vector_avg(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.embedding)).scalar() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_vector_sum(self): sum = Item.select(fn.sum(Item.embedding)).scalar() @@ -177,7 +178,8 @@ def test_vector_sum(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.embedding)).scalar() - assert np.array_equal(sum, np.array([5, 7, 9])) + # does not type cast + assert sum == '[5,7,9]' def test_halfvec_avg(self): avg = Item.select(fn.avg(Item.half_embedding)).scalar() @@ -185,7 +187,8 @@ def test_halfvec_avg(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.half_embedding)).scalar() - assert avg.to_list() == [2.5, 3.5, 4.5] + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_halfvec_sum(self): sum = Item.select(fn.sum(Item.half_embedding)).scalar() @@ -193,7 +196,8 @@ def test_halfvec_sum(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.half_embedding)).scalar() - assert sum.to_list() == [5, 7, 9] + # does not type cast + assert sum == '[5,7,9]' def test_get_or_create(self): Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 24e99d6..9e5100c 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -11,7 +11,7 @@ cur.execute('DROP TABLE IF EXISTS psycopg2_items') cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))') -register_vector(cur) +register_vector(cur, globally=False) class TestPsycopg2: @@ -59,11 +59,11 @@ def test_cursor_factory(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test') cur = conn.cursor(cursor_factory=cursor_factory) - register_vector(cur) + register_vector(cur, globally=False) conn.close() def test_cursor_factory_connection(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) - register_vector(conn) + register_vector(conn, globally=False) conn.close() diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index edce3dc..d7d0756 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -337,7 +337,8 @@ def test_avg(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) avg = session.query(func.avg(Item.embedding)).first()[0] - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_avg_orm(self): with Session(engine) as session: @@ -346,7 +347,8 @@ def test_avg_orm(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) avg = session.scalars(select(func.avg(Item.embedding))).first() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_sum(self): with Session(engine) as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 5685ce6..e1716e1 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -203,7 +203,8 @@ def test_vector_avg(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) avg = session.exec(select(func.avg(Item.embedding))).first() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_vector_sum(self): with Session(engine) as session: @@ -221,7 +222,8 @@ def test_halfvec_avg(self): session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) avg = session.exec(select(func.avg(Item.half_embedding))).first() - assert avg.to_list() == [2.5, 3.5, 4.5] + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_halfvec_sum(self): with Session(engine) as session: