test: test behaviour of cursors with a row maker returning None

diff --git a/tests/test_cursor_common.py b/tests/test_cursor_common.py
index d7bc92c740366702c0f8a01e0b2fb3bfc47696be..2d104daf78a9d1dad4cf9908f2c411ed1c17957b 100644 (file)
--- a/tests/test_cursor_common.py
+++ b/tests/test_cursor_common.py
@@ -893,3 +893,20 @@ def test_error_no_result(conn):
         psycopg.ProgrammingError, match="last operation.*result status: EMPTY_QUERY"
     ):
         cur.fetchone()
+
+
+def test_row_maker_returns_none(conn):
+    cur = conn.cursor(row_factory=rows.scalar_row)
+    query = "values (null), (0)"
+    recs = [None, 0]
+
+    cur.execute(query)
+    assert [cur.fetchone() for _ in range(len(recs))] == recs
+    cur.execute(query)
+    assert cur.fetchmany(len(recs)) == recs
+    cur.execute(query)
+    assert cur.fetchall() == recs
+    cur.execute(query)
+    assert list(cur) == recs
+    stream = cur.stream(query)
+    assert list(stream) == recs

diff --git a/tests/test_cursor_common_async.py b/tests/test_cursor_common_async.py
index 7c7a10e35f75a6f1805ba5f9fea8f87865c05a4a..e219846e65d30927a4690ab2eed9f0cc9bb017a7 100644 (file)
--- a/tests/test_cursor_common_async.py
+++ b/tests/test_cursor_common_async.py
@@ -901,3 +901,20 @@ async def test_error_no_result(aconn):
         psycopg.ProgrammingError, match="last operation.*result status: EMPTY_QUERY"
     ):
         await cur.fetchone()
+
+
+async def test_row_maker_returns_none(aconn):
+    cur = aconn.cursor(row_factory=rows.scalar_row)
+    query = "values (null), (0)"
+    recs = [None, 0]
+
+    await cur.execute(query)
+    assert [await cur.fetchone() for _ in range(len(recs))] == recs
+    await cur.execute(query)
+    assert await cur.fetchmany(len(recs)) == recs
+    await cur.execute(query)
+    assert await cur.fetchall() == recs
+    await cur.execute(query)
+    assert await alist(cur) == recs
+    stream = cur.stream(query)
+    assert await alist(stream) == recs

diff --git a/tests/test_cursor_server.py b/tests/test_cursor_server.py
index fd9192da04e6fcbb771d86cabd87b4700a64803f..01cc8e617320a49678719add65cd623ccfecc4b6 100644 (file)
--- a/tests/test_cursor_server.py
+++ b/tests/test_cursor_server.py
@@ -563,3 +563,20 @@ def test_stolen_cursor_close(conn):
     cur1.execute("declare test cursor for select generate_series(1, 6)")
     cur2 = conn.cursor("test")
     cur2.close()
+
+
+def test_row_maker_returns_none(conn):
+    cur = conn.cursor(row_factory=rows.scalar_row)
+    query = "values (null), (0)"
+    recs = [None, 0]
+
+    cur.execute(query)
+    assert [cur.fetchone() for _ in range(len(recs))] == recs
+    cur.execute(query)
+    assert cur.fetchmany(len(recs)) == recs
+    cur.execute(query)
+    assert cur.fetchall() == recs
+    cur.execute(query)
+    assert list(cur) == recs
+    stream = cur.stream(query)
+    assert list(stream) == recs

diff --git a/tests/test_cursor_server_async.py b/tests/test_cursor_server_async.py
index b94590ae417a37d957990ea5fcad2494124958be..377d057f4bb9aed017cf9a93aeafa1167dbfe726 100644 (file)
--- a/tests/test_cursor_server_async.py
+++ b/tests/test_cursor_server_async.py
@@ -569,3 +569,20 @@ async def test_stolen_cursor_close(aconn):
     await cur1.execute("declare test cursor for select generate_series(1, 6)")
     cur2 = aconn.cursor("test")
     await cur2.close()
+
+
+async def test_row_maker_returns_none(aconn):
+    cur = aconn.cursor(row_factory=rows.scalar_row)
+    query = "values (null), (0)"
+    recs = [None, 0]
+
+    await cur.execute(query)
+    assert [await cur.fetchone() for _ in range(len(recs))] == recs
+    await cur.execute(query)
+    assert await cur.fetchmany(len(recs)) == recs
+    await cur.execute(query)
+    assert await cur.fetchall() == recs
+    await cur.execute(query)
+    assert await alist(cur) == recs
+    stream = cur.stream(query)
+    assert await alist(stream) == recs