Refactor opsqueue_service to use wait_for_server for port assignment and add utility function for server wait logic

.envrc

@@ -11,6 +11,6 @@ if ! has nix_direnv_version || ! nix_direnv_version 3.0.4; then
     source_url "https://raw.githubusercontent.com/nix-community/nix-direnv/3.0.4/direnvrc" "sha256-DzlYZ33mWF/Gs8DDeyjr8mnVmQGx7ASYqA5WlxwvBG4="
 fi
 
-watch_file ./nix/ ./**/*.nix
+watch_file ./nix/ ./**/*.nix default.nix
 
 use nix default.nix --argstr environment shell

default.nix

@@ -8,11 +8,14 @@ let
     p: with p; [
       click
       mypy
-      uv
+      opnieuw
+      psutil
       pytest
-      pytest-random-order
       pytest-parallel
+      pytest-random-order
       pytest-timeout
+      types-psutil
+      uv
 
       # Repeated here so MyPy sees them:
       cbor2

libs/opsqueue_python/tests/conftest.py

@@ -6,6 +6,7 @@
 import subprocess
 import uuid
 import os
+import psutil
 import pytest
 from dataclasses import dataclass
 from pathlib import Path
@@ -14,6 +15,8 @@
 from opsqueue.common import SerializationFormat, json_as_bytes
 from opsqueue.consumer import Strategy
 
+from tests.util import wait_for_server
+
 # @pytest.hookimpl(tryfirst=True)
 # def pytest_configure(config: pytest.Config) -> None:
 #     print("A")
@@ -25,7 +28,7 @@
 @dataclass
 class OpsqueueProcess:
     port: int
-    process: subprocess.Popen[bytes]
+    process: psutil.Popen # subprocess.Popen[bytes]
 
 
 @functools.cache
@@ -52,13 +55,9 @@ def opsqueue() -> Generator[OpsqueueProcess, None, None]:
 
 @contextmanager
 def opsqueue_service(
-    *, port: int | None = None
+    *,
+    port: int = 0,
 ) -> Generator[OpsqueueProcess, None, None]:
-    global test_opsqueue_port_offset
-
-    if port is None:
-        port = random_free_port()
-
     temp_dbname = f"/tmp/opsqueue_tests-{uuid.uuid4()}.db"
 
     command = [
@@ -72,7 +71,8 @@ def opsqueue_service(
     if env.get("RUST_LOG") is None:
         env["RUST_LOG"] = "off"
 
-    with subprocess.Popen(command, cwd=PROJECT_ROOT, env=env) as process:
+    with psutil.Popen(command, cwd=PROJECT_ROOT, env=env) as process:
+        _host, port = wait_for_server(process)
         try:
             wrapper = OpsqueueProcess(port=port, process=process)
             yield wrapper

libs/opsqueue_python/tests/util.py

@@ -0,0 +1,60 @@
+import logging
+
+import psutil
+from opnieuw import retry
+from psutil._common import pconn
+
+LOGGER = logging.getLogger(__name__)
+
+
+@retry(
+    retry_on_exceptions=ValueError,
+    max_calls_total=5,
+    retry_window_after_first_call_in_seconds=5,
+)
+def wait_for_server(proc: psutil.Popen) -> tuple[str, int]:
+    """
+    Wait for a process to be listening on exactly one port and return that address.
+
+    This function expects the given process to listen on a single port only. If the
+    process is listening on no ports, a ValueError is raised and the check is retried.
+    If multiple ports are listening, a RuntimeError is raised as this indicates an
+    unexpected server configuration.
+    """
+    if not proc.is_running():
+        raise ValueError(f"Process {proc} is not running")
+
+    try:
+        # Try to get the connections of the main process first, if that fails try the children.
+        # Processes wrapped with `timeout` do not have connections themselves.
+        connections: list[pconn] = (
+            proc.net_connections()
+            or [
+                child_conn
+                for child in proc.children(recursive=False)
+                for child_conn in child.net_connections()
+            ]
+            or [
+                child_conn
+                for child in proc.children(recursive=True)
+                for child_conn in child.net_connections()
+            ]
+        )
+    except psutil.AccessDenied as e:
+        match proc.status():
+            case psutil.STATUS_ZOMBIE | psutil.STATUS_DEAD | psutil.STATUS_STOPPED:
+                raise RuntimeError(f"Process {proc} has exited unexpectedly") from e
+            case _:
+                raise RuntimeError(
+                    f"Could not get `net_connections` for process {proc}, access denied "
+                ) from e
+
+    ports = [x for x in connections if x.status == psutil.CONN_LISTEN]
+    listen_count = len(ports)
+
+    if listen_count == 0:
+        raise ValueError(f"Process {proc} is not listening on any ports")
+    if listen_count == 1:
+        return ports[0].laddr
+
+    raise RuntimeError(f"Process {proc} is listening on multiple ports")

nix/python-overlay.nix

@@ -1,3 +1,19 @@
 final: prev: {
+  opnieuw = final.buildPythonPackage rec {
+    pname = "opnieuw";
+    version = "3.1.0";
+    pyproject = true;
+
+    propagatedBuildInputs = with final; [
+      setuptools
+      setuptools-scm
+    ];
+
+    src = final.fetchPypi {
+      inherit pname version;
+      sha256 = "sha256-4aLFPQYqnCOBsxdXr41gr5sr/cT9HcN9PLEc+AWwLrY=";
+    };
+  };
+
   opsqueue_python = final.callPackage ../libs/opsqueue_python/opsqueue_python.nix { };
 }