test: add test to verify stream+close interaction

diff --git a/tests/test_cursor.py b/tests/test_cursor.py
index 9f02710c9a11b26feb5fe0831a4bfcf204b1fb87..acd79c08bfff0d6f36cf9c776b5e922e08a89cd9 100644 (file)
--- a/tests/test_cursor.py
+++ b/tests/test_cursor.py
@@ -587,6 +587,18 @@ def test_stream_error_notx(conn):
     assert conn.info.transaction_status == conn.TransactionStatus.IDLE
 
 
+def test_stream_close(conn):
+    cur = conn.cursor()
+    with pytest.raises(psycopg.OperationalError):
+        for rec in cur.stream("select generate_series(1, 3)"):
+            if rec[0] == 1:
+                conn.close()
+            else:
+                assert False
+
+    assert conn.closed
+
+
 def test_stream_binary_cursor(conn):
     cur = conn.cursor(binary=True)
     recs = []

diff --git a/tests/test_cursor_async.py b/tests/test_cursor_async.py
index 29c1a583a7409538ed81cc32719762d29a97e6e3..2bd029c0134c56c95fabd25a4d63bd10198fc552 100644 (file)
--- a/tests/test_cursor_async.py
+++ b/tests/test_cursor_async.py
@@ -579,6 +579,19 @@ async def test_stream_error_notx(aconn):
     assert aconn.info.transaction_status == aconn.TransactionStatus.IDLE
 
 
+async def test_stream_close(aconn):
+    await aconn.set_autocommit(True)
+    cur = aconn.cursor()
+    with pytest.raises(psycopg.OperationalError):
+        async for rec in cur.stream("select generate_series(1, 3)"):
+            if rec[0] == 1:
+                await aconn.close()
+            else:
+                assert False
+
+    assert aconn.closed
+
+
 async def test_stream_binary_cursor(aconn):
     cur = aconn.cursor(binary=True)
     recs = []