Skip to content

Commit 057eff2

Browse files
committed
Improved tests [skip ci]
1 parent e566d4c commit 057eff2

File tree

2 files changed

+3
-6
lines changed

2 files changed

+3
-6
lines changed

tests/test_psycopg2.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,17 +82,15 @@ def test_halfvec_array(self):
8282

8383
cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id')
8484
res = cur.fetchone()
85-
assert res[0][0] == HalfVector([1.5, 2, 3])
86-
assert res[0][1] == HalfVector([4.5, 5, 6])
85+
assert res[0] == [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])]
8786

8887
def test_sparsevec_array(self):
8988
embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])]
9089
cur.execute('INSERT INTO psycopg2_items (sparse_embeddings) VALUES (%s::sparsevec[])', (embeddings,))
9190

9291
cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id')
9392
res = cur.fetchone()
94-
assert res[0][0] == SparseVector([1.5, 2, 3])
95-
assert res[0][1] == SparseVector([4.5, 5, 6])
93+
assert res[0] == [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])]
9694

9795
def test_cursor_factory(self):
9896
for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]:

tests/test_sqlalchemy.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -551,8 +551,7 @@ def test_halfvec_array(self, engine):
551551

552552
# this fails if the driver does not cast arrays
553553
item = session.get(Item, 1)
554-
assert item.half_embeddings[0] == HalfVector([1, 2, 3])
555-
assert item.half_embeddings[1] == HalfVector([4, 5, 6])
554+
assert item.half_embeddings == [HalfVector([1, 2, 3]), HalfVector([4, 5, 6])]
556555

557556

558557
@pytest.mark.parametrize('engine', async_engines)

0 commit comments

Comments
 (0)