tests: simplify PGcancelConn tests

We split the previous (misnamed) test_cancel_nonblocking() into
test_cancel_conn_blocking() and test_cancel_conn_nonblocking()
respectively checking PQcancelBlocking() and
PQcancelStart()/PQcancelPoll().

The cancellable_query() is now responsible for creating the monitor_conn
and uses a PGconn value, coming from the pgconn fixture in tests.

diff --git a/tests/pq/test_pgconn.py b/tests/pq/test_pgconn.py
index 50877929fa0a7f478e73d6ab6e18d43b63ec020f..2c66b098e7fd64e1293a625d4f7584b8c161c588 100644 (file)
--- a/tests/pq/test_pgconn.py
+++ b/tests/pq/test_pgconn.py
@@ -386,8 +386,15 @@ def test_set_single_row_mode(pgconn):
 
 
 @contextlib.contextmanager
-def cancellable_query(pgconn: PGconn, monitor_conn: PGconn) -> Iterator[None]:
+def cancellable_query(pgconn: PGconn) -> Iterator[None]:
+    dsn = b" ".join(b"%s='%s'" % (i.keyword, i.val) for i in pgconn.info if i.val)
+    monitor_conn = pq.PGconn.connect(dsn)
+    assert (
+        monitor_conn.status == pq.ConnStatus.OK
+    ), f"bad connection: {monitor_conn.error_message.decode('utf8', 'replace')}"
+
     pgconn.send_query_params(b"SELECT pg_sleep($1)", [b"180"])
+
     while True:
         r = monitor_conn.exec_(
             b"SELECT count(*) FROM pg_stat_activity"
@@ -396,6 +403,7 @@ def cancellable_query(pgconn: PGconn, monitor_conn: PGconn) -> Iterator[None]:
         )
         assert r.status == pq.ExecStatus.TUPLES_OK
         if r.get_value(0, 0) != b"0":
+            del monitor_conn
             break
 
         time.sleep(0.01)
@@ -411,44 +419,36 @@ def cancellable_query(pgconn: PGconn, monitor_conn: PGconn) -> Iterator[None]:
 
 
 @pytest.mark.libpq(">= 17")
-def test_cancel_nonblocking(dsn):
-    # mimic test_cancel() from src/test/modules/libpq_pipeline/libpq_pipeline.c
-    def connect() -> PGconn:
-        conn = pq.PGconn.connect(dsn.encode())
-        if conn.status != pq.ConnStatus.OK:
-            pytest.fail(
-                f"bad connection: {conn.error_message.decode('utf8', 'replace')}"
-            )
-        return conn
-
-    conn, monitor_conn = connect(), connect()
-    conn.nonblocking = 1
+def test_cancel_conn_blocking(pgconn):
+    # test PQcancelBlocking, similarly to test_cancel() from
+    # src/test/modules/libpq_pipeline/libpq_pipeline.c
+    pgconn.nonblocking = 1
 
-    # test PQcancel
-    with cancellable_query(conn, monitor_conn):
-        cancel = conn.get_cancel()
-        cancel.cancel()
+    with cancellable_query(pgconn):
+        cancel_conn = pgconn.cancel_conn()
+        assert cancel_conn.status == pq.ConnStatus.ALLOCATED
+        cancel_conn.blocking()
+        assert cancel_conn.status == pq.ConnStatus.OK
 
-    # PGcancel object can be reused for the next query
-    with cancellable_query(conn, monitor_conn):
-        cancel.cancel()
+    # test PQcancelReset works on the cancel connection and it can be reused
+    # after
+    cancel_conn.reset()
+    with cancellable_query(pgconn):
+        cancel_conn.blocking()
+        assert cancel_conn.status == pq.ConnStatus.OK
 
-    del cancel
 
-    # test PQcancelBlocking
-    with cancellable_query(conn, monitor_conn):
-        cancel_conn = conn.cancel_conn()
-        assert cancel_conn.status == pq.ConnStatus.ALLOCATED
-        cancel_conn.blocking()
-        assert cancel_conn.status == pq.ConnStatus.OK  # type: ignore[comparison-overlap] # noqa: E501
-    cancel_conn.finish()
-    del cancel_conn
+@pytest.mark.libpq(">= 17")
+def test_cancel_conn_nonblocking(pgconn):
+    # test PQcancelStart() and then polling with PQcancelPoll, similarly to
+    # test_cancel() from src/test/modules/libpq_pipeline/libpq_pipeline.c
+    pgconn.nonblocking = 1
 
     wait_cancel = partial(wait, poll_method="poll", timeout=3)
 
     # test PQcancelCreate and then polling with PQcancelPoll
-    with cancellable_query(conn, monitor_conn):
-        cancel_conn = conn.cancel_conn()
+    with cancellable_query(pgconn):
+        cancel_conn = pgconn.cancel_conn()
         assert cancel_conn.status == pq.ConnStatus.ALLOCATED
         cancel_conn.start()
         assert cancel_conn.status == pq.ConnStatus.STARTED
@@ -458,13 +458,11 @@ def test_cancel_nonblocking(dsn):
     # test PQcancelReset works on the cancel connection and it can be reused
     # after
     cancel_conn.reset()
-    with cancellable_query(conn, monitor_conn):
+    with cancellable_query(pgconn):
         cancel_conn.start()
         wait_cancel(cancel_conn)
         assert cancel_conn.status == pq.ConnStatus.OK
 
-    cancel_conn.finish()
-
 
 def test_cancel(pgconn):
     cancel = pgconn.get_cancel()