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 2ed345de49fea9887317bd821938b6c8d9d0aa6f..12a5f93e77eff92eb1f17c704486b24911007482 100644 (file)
--- 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 36f80a1c0a55ddae867c7e0ccd4fe6afa690ad31..7abfc58a964555b9e8c0a85094bdeac9b95bb97a 100644 (file)
--- 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):