Added tests for copy

ankane · ankane · commit f2d5dab2e5a6 · 2023-05-03T19:18:05.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,6 @@
 ## 0.1.7 (unreleased)
 
-- Fixed binary copy for psycopg3
+- Fixed `set_types` for psycopg3
 
 ## 0.1.6 (2022-05-22)
 
diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py
@@ -51,3 +51,22 @@ def test_binary_format_non_contiguous(self):
         assert not embedding.data.contiguous
         res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0]
         assert np.array_equal(res, np.array([3, 2, 1.5]))
+
+    def test_text_copy(self):
+        embedding = np.array([1.5, 2, 3])
+        cur = conn.cursor()
+        with cur.copy("COPY item (embedding) FROM STDIN") as copy:
+            copy.write_row([embedding])
+
+    def test_binary_copy(self):
+        embedding = np.array([1.5, 2, 3])
+        cur = conn.cursor()
+        with cur.copy("COPY item (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy:
+            copy.write_row([embedding])
+
+    def test_binary_copy_set_types(self):
+        embedding = np.array([1.5, 2, 3])
+        cur = conn.cursor()
+        with cur.copy("COPY item (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy:
+            copy.set_types(['int8', 'vector'])
+            copy.write_row([1, embedding])
