From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 1 Nov 2022 22:39:22 +0000 (+0100)
Subject: test: allow choosing a wait function using the PSYCOPG_WAIT_FUNC env var
X-Git-Tag: 3.1.5~7^2~4
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=459912137fa069871fa9b901a80887cc29787386;p=thirdparty%2Fpsycopg.git

test: allow choosing a wait function using the PSYCOPG_WAIT_FUNC env var

We won't make this public yet, possibly we never will (I think we are
the people responsible for this choice), but it might be useful for
testing.
---

diff --git a/psycopg/psycopg/abc.py b/psycopg/psycopg/abc.py
index 2ed345de4..12a5f93e7 100644
--- a/psycopg/psycopg/abc.py
+++ b/psycopg/psycopg/abc.py
@@ -46,6 +46,17 @@ PQGen: TypeAlias = Generator["Wait", "Ready", RV]
 """
 
 
+class WaitFunc(Protocol):
+    """
+    Wait on the connection which generated `PQgen` and return its final result.
+    """
+
+    def __call__(
+        self, gen: PQGen[RV], fileno: int, timeout: Optional[float] = None
+    ) -> RV:
+        ...
+
+
 # Adaptation types
 
 DumpFunc: TypeAlias = Callable[[Any], Buffer]
diff --git a/psycopg/psycopg/waiting.py b/psycopg/psycopg/waiting.py
index 36f80a1c0..7abfc58a9 100644
--- a/psycopg/psycopg/waiting.py
+++ b/psycopg/psycopg/waiting.py
@@ -9,6 +9,7 @@ These functions are designed to consume the generators returned by the
 # Copyright (C) 2020 The Psycopg Team
 
 
+import os
 import select
 import selectors
 from typing import Dict, Optional
@@ -16,7 +17,7 @@ from asyncio import get_event_loop, wait_for, Event, TimeoutError
 from selectors import DefaultSelector
 
 from . import errors as e
-from .abc import PQGen, PQGenConn, RV
+from .abc import RV, PQGen, PQGenConn, WaitFunc
 from ._enums import Wait as Wait, Ready as Ready  # re-exported
 from ._cmodule import _psycopg
 
@@ -297,7 +298,19 @@ if _psycopg:
 # the selectors objects have a generic interface but come with some overhead,
 # so we also offer more finely tuned implementations.
 
-if _psycopg:
+wait: WaitFunc
+
+# Allow the user to choose a specific function for testing
+if "PSYCOPG_WAIT_FUNC" in os.environ:
+    fname = os.environ["PSYCOPG_WAIT_FUNC"]
+    if not fname.startswith("wait_") or fname not in globals():
+        raise ImportError(
+            "PSYCOPG_WAIT_FUNC should be the name of an available wait function;"
+            f" got {fname!r}"
+        )
+    wait = globals()[fname]
+
+elif _psycopg:
     wait = wait_c
 
 elif selectors.DefaultSelector is getattr(selectors, "SelectSelector", None):