From 4a7ea57b3b0045fec92d8cee40a16211ced9d499 Mon Sep 17 00:00:00 2001
From: Denis Laxalde <denis.laxalde@dalibo.com>
Date: Mon, 6 Dec 2021 08:44:02 +0100
Subject: [PATCH] Add a --pq-tracefile pytest option

When specified, we'd enable tracing of the client/server communication
to a file. The file is opened for the whole test session and we set a
title section before each test run.
---
 tests/fix_db.py | 66 ++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 60 insertions(+), 6 deletions(-)

diff --git a/tests/fix_db.py b/tests/fix_db.py
index 3e4c60ab8..1797177de 100644
--- a/tests/fix_db.py
+++ b/tests/fix_db.py
@@ -1,8 +1,11 @@
 import os
 import pytest
 import logging
+from contextlib import contextmanager
 from typing import List
 
+import psycopg
+from psycopg import pq
 from .utils import check_server_version
 
 
@@ -14,6 +17,12 @@ def pytest_addoption(parser):
         help="Connection string to run database tests requiring a connection"
         " [you can also use the PSYCOPG_TEST_DSN env var].",
     )
+    parser.addoption(
+        "--pq-tracefile",
+        metavar="TRACEFILE",
+        default=None,
+        help="Generate a libpq trace to TRACEFILE.",
+    )
 
 
 def pytest_configure(config):
@@ -51,8 +60,48 @@ def dsn(request):
     return dsn
 
 
+@pytest.fixture(scope="session")
+def tracefile(request):
+    """Open and yield a file for libpq client/server communication traces if
+    --pq-tracefile option is set.
+    """
+    tracefile = request.config.getoption("--pq-tracefile")
+    if not tracefile:
+        yield None
+        return
+
+    with open(tracefile, "w") as f:
+        yield f
+
+
+@contextmanager
+def maybe_trace(pgconn, tracefile, function):
+    """Handle libpq client/server communication traces for a single test
+    function.
+    """
+    if tracefile is None:
+        yield None
+        return
+
+    title = f" {function.__module__}::{function.__qualname__} ".center(80, "=")
+    tracefile.write(title + "\n")
+    tracefile.flush()
+
+    pgconn.trace(tracefile.fileno())
+    try:
+        pgconn.set_trace_flags(
+            pq.Trace.SUPPRESS_TIMESTAMPS | pq.Trace.REGRESS_MODE
+        )
+    except psycopg.NotSupportedError:
+        pass
+    try:
+        yield None
+    finally:
+        pgconn.untrace()
+
+
 @pytest.fixture
-def pgconn(dsn, request):
+def pgconn(dsn, request, tracefile):
     """Return a PGconn connection open to `--test-dsn`."""
     from psycopg import pq
 
@@ -65,12 +114,15 @@ def pgconn(dsn, request):
     if msg:
         conn.finish()
         pytest.skip(msg)
-    yield conn
+
+    with maybe_trace(conn, tracefile, request.function):
+        yield conn
+
     conn.finish()
 
 
 @pytest.fixture
-def conn(dsn, request):
+def conn(dsn, request, tracefile):
     """Return a `Connection` connected to the ``--test-dsn`` database."""
     from psycopg import Connection
 
@@ -79,12 +131,13 @@ def conn(dsn, request):
     if msg:
         conn.close()
         pytest.skip(msg)
-    yield conn
+    with maybe_trace(conn.pgconn, tracefile, request.function):
+        yield conn
     conn.close()
 
 
 @pytest.fixture
-async def aconn(dsn, request):
+async def aconn(dsn, request, tracefile):
     """Return an `AsyncConnection` connected to the ``--test-dsn`` database."""
     from psycopg import AsyncConnection
 
@@ -93,7 +146,8 @@ async def aconn(dsn, request):
     if msg:
         await conn.close()
         pytest.skip(msg)
-    yield conn
+    with maybe_trace(conn.pgconn, tracefile, request.function):
+        yield conn
     await conn.close()
 
 
-- 
2.47.2