trainUs/tests/conftest.py

import urllib.request
import pytest
import subprocess
import time
import signal
import os
import fcntl


def _counter_add(delta: int, count_file: str, lock_file: str) -> int:
    """Atomically add `delta` to the worker counter and return the new value."""
    with open(lock_file, "w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        try:
            try:
                with open(count_file) as f:
                    value = int(f.read().strip() or "0")
            except FileNotFoundError:
                value = 0
            value += delta
            with open(count_file, "w") as f:
                f.write(str(value))
            return value
        finally:
            fcntl.flock(lock, fcntl.LOCK_UN)


def _counter_read(count_file: str, lock_file: str) -> int:
    """Read the current worker counter value without modifying it."""
    with open(lock_file, "w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        try:
            try:
                with open(count_file) as f:
                    return int(f.read().strip() or "0")
            except FileNotFoundError:
                return 0
        finally:
            fcntl.flock(lock, fcntl.LOCK_UN)


@pytest.fixture(scope="session")
def app_url(tmp_path_factory, worker_id):
    """Start the Vite dev server and return its URL.

    When running under pytest-xdist, only the first worker (gw0) starts the
    server and writes the URL to a shared temp file.  All other workers wait
    for that file to appear and then read the URL from it.  The server is only
    stopped once every worker has finished (tracked via a shared counter file).
    When running without xdist (worker_id == "master") the behaviour is
    identical to the original single-worker case.

    Shared files are placed inside tmp_path_factory.getbasetemp() so that two
    concurrent pytest invocations never collide (each invocation gets a unique
    basetemp directory from pytest's per-run counter).
    """
    basetemp = tmp_path_factory.getbasetemp()
    # In xdist each worker gets its own subdirectory; step up to the shared parent
    if worker_id != "master":
        basetemp = basetemp.parent
    server_url_file = str(basetemp / "trainus_server_url")
    worker_count_file = str(basetemp / "trainus_worker_count")
    worker_count_lock = str(basetemp / "trainus_worker_count.lock")

    url = "http://localhost:5199"
    is_controller = worker_id in ("master", "gw0")

    # Register this worker in the shared counter
    _counter_add(1, worker_count_file, worker_count_lock)

    server_error_file = str(basetemp / "trainus_server_error")

    if is_controller:
        # Remove stale files from a previous run if present
        for stale in (server_url_file, server_error_file):
            try:
                os.remove(stale)
            except FileNotFoundError:
                pass

        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        process = subprocess.Popen(
            ["npm", "run", "dev", "--", "--port", "5199", "--strictPort"],
            cwd=project_root,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            preexec_fn=os.setsid,
        )

        # Wait for the server to be ready (up to 60 s)
        for _ in range(120):
            try:
                urllib.request.urlopen(url, timeout=1)
                break
            except Exception:
                time.sleep(0.5)
        else:
            process.kill()
            # Signal secondary workers to stop waiting
            with open(server_error_file, "w") as f:
                f.write("Vite dev server failed to start within 60s")
            raise RuntimeError("Vite dev server failed to start within 60s")

        # Publish the URL so other workers can pick it up
        with open(server_url_file, "w") as f:
            f.write(url)

        yield url

        # Mark the controller itself as done, then wait for all other workers
        # to finish before killing the server.  We decrement exactly once here
        # and then only *read* the counter so we never accidentally modify it
        # inside the polling loop.
        _counter_add(-1, worker_count_file, worker_count_lock)
        for _ in range(240):
            if _counter_read(worker_count_file, worker_count_lock) <= 0:
                break
            time.sleep(0.5)

        # Cleanup
        os.killpg(os.getpgid(process.pid), signal.SIGTERM)

    else:
        # Secondary workers: wait until gw0 has written the URL file (up to 90 s)
        for _ in range(180):
            if os.path.exists(server_url_file):
                break
            if os.path.exists(server_error_file):
                with open(server_error_file) as f:
                    msg = f.read().strip()
                raise RuntimeError(f"gw0 failed to start the Vite dev server: {msg}")
            time.sleep(0.5)
        else:
            raise RuntimeError("Timed out waiting for Vite dev server URL from gw0")

        with open(server_url_file) as f:
            server_url = f.read().strip()

        yield server_url

        # Signal this worker is done
        _counter_add(-1, worker_count_file, worker_count_lock)


@pytest.fixture(scope="session")
def browser_type_launch_args():
    """Configure browser launch args."""
    return {}


@pytest.fixture(scope="session")
def browser_context_args():
    """Configure browser context."""
    return {
        "viewport": {"width": 1280, "height": 720},
    }