Add tests for cursor's row_factory with nextset()

diff --git a/tests/test_cursor.py b/tests/test_cursor.py
index cdd9f6de1901d10d0b1d7249703a1788c66f8d88..9db2ab3581935eaddcd56923c5e795d11dcb8bc0 100644 (file)
--- a/tests/test_cursor.py
+++ b/tests/test_cursor.py
@@ -272,6 +272,12 @@ def test_row_factory(conn):
     r = cur.fetchall()
     assert r == [[-1], [-2], [-3]]
 
+    cur.execute("select 42; select generate_series(1,3)")
+    assert cur.fetchall() == [[-42]]
+    assert cur.nextset()
+    assert cur.fetchall() == [[-1], [-2], [-3]]
+    assert cur.nextset() is None
+
 
 def test_query_params_execute(conn):
     cur = conn.cursor()

diff --git a/tests/test_cursor_async.py b/tests/test_cursor_async.py
index 1aa4f219837ea685f6a9d1fc7df4c7dda161e5f5..fba7fb502b8a98a53adc835b74f091341cec4e75 100644 (file)
--- a/tests/test_cursor_async.py
+++ b/tests/test_cursor_async.py
@@ -270,10 +270,9 @@ async def test_iter_stop(aconn):
 
 async def test_row_factory(aconn):
     def my_row_factory(cursor):
-        assert cursor.description is not None
-        titles = [c.name for c in cursor.description]
-
         def mkrow(values):
+            assert cursor.description is not None
+            titles = [c.name for c in cursor.description]
             return [
                 f"{value.upper()}{title}"
                 for title, value in zip(titles, values)
@@ -286,6 +285,12 @@ async def test_row_factory(aconn):
     (r,) = await cur.fetchone()
     assert r == "FOObar"
 
+    await cur.execute("select 'x' as x; select 'y' as y, 'z' as z")
+    assert await cur.fetchall() == [["Xx"]]
+    assert cur.nextset()
+    assert await cur.fetchall() == [["Yy", "Zz"]]
+    assert cur.nextset() is None
+
 
 async def test_query_params_execute(aconn):
     cur = await aconn.cursor()