From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 7 Jun 2022 07:41:14 +0000 (+0200)
Subject: test: add test to verify stream+close interaction
X-Git-Tag: 3.1~66^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7a44e94c6f8afb4948217d025e2d27348e36fba8;p=thirdparty%2Fpsycopg.git

test: add test to verify stream+close interaction
---

diff --git a/tests/test_cursor.py b/tests/test_cursor.py
index 8e7c78500..7c7db7b95 100644
--- 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 3fcbea0ee..232f1ae42 100644
--- 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 = []