Improved test

pgvector · Nov 5, 2024 · 706cebc · 706cebc
1 parent 75e14d8
commit 706cebc
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py
@@ -101,11 +101,14 @@ def test_pool(self):
 
         conn = pool.getconn()
         try:
-            cur = conn.cursor()
-
             # use globally=True for apps
-            register_vector(cur, globally=False)
+            register_vector(conn, globally=False)
+        finally:
+            pool.putconn(conn)
 
+        conn = pool.getconn()
+        try:
+            cur = conn.cursor()
             cur.execute("SELECT '[1,2,3]'::vector")
             res = cur.fetchone()
             assert np.array_equal(res[0], np.array([1, 2, 3]))