test: add test to verify stream+close interaction

diff --git a/tests/test_cursor.py b/tests/test_cursor.py
index 8e7c785008fc946207f5ac92ed1d325d37f50d9e..7c7db7b95208a536588101d24998b1d760909454 100644 (file)
--- a/tests/test_cursor.py
+++ b/tests/test_cursor.py
@@ -639,6 +639,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 3fcbea0ee7154ec19b2ef910d4aaf94fb46ebe41..232f1ae4214b386a9740f12f9e552868ccd1689e 100644 (file)
--- a/tests/test_cursor_async.py
+++ b/tests/test_cursor_async.py
@@ -631,6 +631,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 = []