Skip to content

Commit f2f30e1

Browse files
committed
Improved test for arrays with Django [skip ci]
1 parent 417f3ec commit f2f30e1

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

tests/test_django.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -440,5 +440,11 @@ def test_missing(self):
440440
def test_vector_array(self):
441441
Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]).save()
442442

443-
# this fails if the driver does not cast arrays
444-
# item = Item.objects.get(pk=1)
443+
with connection.cursor() as cursor:
444+
from pgvector.psycopg import register_vector
445+
register_vector(cursor.connection)
446+
447+
# this fails if the driver does not cast arrays
448+
item = Item.objects.get(pk=1)
449+
assert item.embeddings[0].tolist() == [1, 2, 3]
450+
assert item.embeddings[1].tolist() == [4, 5, 6]

0 commit comments

Comments
 (0)