From a8a1bf106c079d330e1af3d9bbfb02df84fa6de6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 25 Aug 2024 12:38:22 -0700 Subject: [PATCH] Added tests for Psycopg 2 cursor factories --- tests/test_psycopg2.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 54da6a7..e6ba996 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,6 +1,7 @@ import numpy as np from pgvector.psycopg2 import register_vector, SparseVector import psycopg2 +from psycopg2.extras import DictCursor, NamedTupleCursor conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -53,3 +54,16 @@ def test_sparsevec(self): res = cur.fetchall() assert res[0][0].to_list() == [1.5, 2, 3] assert res[1][0] is None + + def test_cursor_factory(self): + for cursor_factory in [DictCursor, NamedTupleCursor]: + conn = psycopg2.connect(dbname='pgvector_python_test') + cur = conn.cursor(cursor_factory=cursor_factory) + register_vector(cur) + conn.close() + + def test_cursor_factory_connection(self): + for cursor_factory in [DictCursor, NamedTupleCursor]: + conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) + register_vector(conn) + conn.close()