litellm/tests/integration/database/test_harness_lifecycle.py

from __future__ import annotations

import asyncio
import errno
import os
import re
import signal
import socket
import subprocess
import sys
import time
from collections import Counter
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from contextlib import AbstractContextManager, ExitStack
from pathlib import Path
from queue import SimpleQueue
from threading import Event
from tempfile import TemporaryDirectory
from types import FrameType
from typing import Final, Literal

import httpx
import psutil
import psycopg
import pytest
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict

from tests.integration._support import process as process_support
from tests.integration._support.async_server import LoopbackServer
from tests.integration._support.client import Gateway, eventually
from tests.integration._support.database_relay import (
    DatabaseRelay,
    DroppedConnectionRelay,
    HeldStatementRelay,
    database_relay,
    dropped_connection_relay,
    held_statement_relay,
)
from tests.integration._support.runtime import configure_executor
from tests.integration._support.sharding import ShardExecution, shard_errors


class ExitRecord(TypedDict):
    root_pid: ReadOnly[int]
    exit_after_stop: ReadOnly[int | None]
    exception_at_record: ReadOnly[str | None]


Relay = DatabaseRelay | DroppedConnectionRelay | HeldStatementRelay
RelayFactory = Callable[[str, bytes], AbstractContextManager[tuple[Relay, str]]]
STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")


@pytest.mark.parametrize("factory", (database_relay, dropped_connection_relay, held_statement_relay))
def test_relay_closes_active_connections_and_releases_its_listener(factory: RelayFactory) -> None:
    with ExitStack() as cleanup:
        with factory(os.environ["DATABASE_URL"], b"unissued-lifecycle-trigger") as (relay, url):
            connection: Final = cleanup.enter_context(psycopg.connect(url, connect_timeout=2, autocommit=True))
            assert connection.execute("SELECT 1").fetchone() == (1,)
        with pytest.raises(psycopg.OperationalError):
            connection.execute("SELECT 1")
        with socket.socket() as replacement:
            replacement.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            replacement.bind(("127.0.0.1", relay.port))
            replacement.listen()
        assert relay._loop.is_closed()


async def _unused_peer(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    await reader.read()


def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> None:
    with socket.socket() as occupied:
        occupied.bind(("127.0.0.1", 0))
        occupied.listen()
        relay: Final = LoopbackServer(occupied.getsockname()[1], _unused_peer)
        with pytest.raises(OSError, match=os.strerror(errno.EADDRINUSE).lower()):
            relay.start()
        relay.stop()
        assert relay.loop.is_closed()


def test_upstream_readiness_failure_reaps_the_started_process(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    monkeypatch.setattr(process_support, "_UPSTREAM_READY_SECONDS", 0)
    children_before: Final = frozenset(child.pid for child in psutil.Process().children(recursive=True))
    slot: Final = process_support.UpstreamSlot(tmp_path, process_support._free_port(), process_support._harness_root())
    try:
        with pytest.raises(AssertionError, match="readiness deadline exceeded"):
            slot.start()
        assert slot.process is None
        assert frozenset(child.pid for child in psutil.Process().children(recursive=True)) == children_before
    finally:
        if slot.process is not None:
            slot.stop()


class _InterruptedCleanup(BaseException):
    pass


async def test_owned_runtime_starts_the_entire_held_provider_burst_on_a_small_runner() -> None:
    arrived: Final[SimpleQueue[bool]] = SimpleQueue()
    complete: Final = Event()
    release: Final = Event()

    def held() -> None:
        arrived.put(True)
        if arrived.qsize() == 20:
            complete.set()
        assert release.wait(timeout=5), "Held executor job was never released"

    with ThreadPoolExecutor(max_workers=1) as initial:
        asyncio.get_running_loop().set_default_executor(initial)
        configure_executor()
        async with asyncio.TaskGroup() as tasks:
            burst: Final = tuple(tasks.create_task(asyncio.to_thread(held)) for _ in range(20))
            try:
                assert await asyncio.wait_for(asyncio.to_thread(complete.wait, 2), timeout=3)
                assert arrived.qsize() == len(burst)
            finally:
                release.set()


@pytest.mark.timeout(90)
def test_slow_worker_startup_survives_healthchecks_and_a_stalled_worker_is_replaced(
    gateway: Gateway, tmp_path: Path
) -> None:
    hook: Final = tmp_path / "slow_worker_startup.py"
    hook.write_text("from threading import Event\n\ndef configure():\n    Event().wait(6)\n")
    overrides: Final = {
        "PYTHONPATH": os.pathsep.join((str(tmp_path), os.environ.get("PYTHONPATH", ""))),
        "LITELLM_WORKER_STARTUP_HOOKS": "slow_worker_startup:configure",
    }
    with process_support.owned_proxy_process(gateway, tmp_path, overrides, workers=2) as owned:
        workers: Final = eventually(
            lambda: tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text())),
            lambda pids: len(pids) == 2 and owned.log.read_text().count("Application startup complete.") == 2,
            seconds=30,
        )
        victim: Final = psutil.Process(workers[0])
        victim.suspend()
        try:
            replacement: Final = eventually(
                lambda: tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text())),
                lambda pids: len(pids) == 3 and owned.log.read_text().count("Application startup complete.") == 3,
                seconds=30,
            )
            assert replacement[-1] not in workers
            assert not psutil.pid_exists(victim.pid), "Supervisor did not reap the stalled worker"
            response: Final = owned.gateway.request("GET", "/health/readiness")
            assert response.status_code == 200, response.text
        finally:
            try:
                victim.kill()
            except psutil.NoSuchProcess:
                pass


def test_interrupted_cleanup_reaps_the_owned_root_and_its_child() -> None:
    child_code: Final = (
        "import signal,time; signal.signal(signal.SIGTERM,signal.SIG_IGN); print('ready',flush=True); time.sleep(60)"
    )
    parent_code: Final = (
        "import signal,subprocess,sys,time; signal.signal(signal.SIGTERM,signal.SIG_IGN); "
        f"child=subprocess.Popen((sys.executable,'-P','-c',{child_code!r}),stdout=subprocess.PIPE); "
        "assert child.stdout.readline()==b'ready\\n'; print('ready',flush=True); time.sleep(60)"
    )

    def interrupt(signum: int, frame: FrameType | None) -> None:
        raise _InterruptedCleanup

    with subprocess.Popen(
        (sys.executable, "-P", "-c", parent_code), start_new_session=True, stdout=subprocess.PIPE
    ) as process:
        try:
            assert process.stdout is not None
            assert process.stdout.readline() == b"ready\n"
            assert len(process_support.group_members(process.pid)) == 2
            previous_timer: Final = signal.getitimer(signal.ITIMER_REAL)
            started: Final = time.monotonic()
            previous_handler: Final = signal.signal(signal.SIGALRM, interrupt)
            try:
                signal.setitimer(signal.ITIMER_REAL, 0.2)
                with pytest.raises(_InterruptedCleanup):
                    process_support._stop(process)
            finally:
                signal.setitimer(signal.ITIMER_REAL, 0)
                signal.signal(signal.SIGALRM, previous_handler)
                remaining: Final = (
                    max(0.001, previous_timer[0] - (time.monotonic() - started)) if previous_timer[0] else 0
                )
                signal.setitimer(signal.ITIMER_REAL, remaining, previous_timer[1])
            assert process.poll() is not None
            eventually(lambda: process_support.group_members(process.pid), lambda members: not members, seconds=5)
        finally:
            process_support.signal_group(process.pid, signal.SIGKILL)
            process.wait(timeout=3)


@pytest.mark.parametrize("explicit_dotenv", (False, True))
def test_owned_process_ignores_ambient_dotenv_unless_its_contract_opts_in(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path, explicit_dotenv: bool
) -> None:
    monkeypatch.delenv("OWNED_DOTENV_SENTINEL", raising=False)
    monkeypatch.delenv("PYTHON_DOTENV_DISABLED", raising=False)
    (tmp_path / ".env").write_text("OWNED_DOTENV_SENTINEL=ambient-config\n")
    with httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client:
        environment: Final = process_support._proxy_environment(
            Gateway(client, "owned-test-key", "http://127.0.0.1:1"),
            {"PYTHON_DOTENV_DISABLED": "0"} if explicit_dotenv else {},
            (),
        )
    result: Final = subprocess.run(
        (
            sys.executable,
            "-P",
            "-c",
            "import os; from prisma import Prisma; Prisma(); print(os.getenv('OWNED_DOTENV_SENTINEL', 'absent'))",
        ),
        cwd=tmp_path,
        env=dict(environment),
        capture_output=True,
        text=True,
        timeout=10,
        check=True,
    )
    assert result.stdout.strip() == ("ambient-config" if explicit_dotenv else "absent")


def test_seeded_collection_reuses_scoped_fixtures_and_executes_every_case(tmp_path: Path) -> None:
    events: Final = tmp_path / "events.txt"
    probe: Final = tmp_path / "test_fixture_probe.py"
    probe.write_text(
        "import pytest\nfrom pathlib import Path\n"
        f"EVENTS = Path({str(events)!r})\n"
        "CONFIGS = ('D1',)*5 + ('D2',)*5 + ('D3',)*4 + ('D4',)*4\n"
        "def record(value):\n"
        "    with EVENTS.open('a') as output:\n"
        "        output.write(value + '\\n')\n"
        "@pytest.fixture(scope='module')\n"
        "def rig(request):\n"
        "    record('start ' + request.param)\n"
        "    yield request.param\n"
        "    record('stop ' + request.param)\n"
        "@pytest.mark.parametrize(('rig', 'number'), "
        "tuple((value, index) for index, value in enumerate(CONFIGS)), indirect=['rig'], scope='module')\n"
        "def test_route(rig, number):\n"
        "    assert rig == CONFIGS[number]\n"
        "    record('case ' + str(number))\n"
    )
    result: Final = subprocess.run(
        (
            sys.executable,
            "-P",
            "-m",
            "pytest",
            str(probe),
            "-p",
            "tests.integration.conftest",
            "--integration-order-seed=1848224111",
            "--timeout=15",
            "-q",
        ),
        env={
            name: value
            for name, value in os.environ.items()
            if name not in ("INTEGRATION_RESULTS_DIR", "INTEGRATION_ROUTING")
        },
        capture_output=True,
        text=True,
        timeout=30,
        check=False,
    )
    assert result.returncode == 0, result.stdout + result.stderr
    observed: Final = tuple(events.read_text().splitlines())
    configurations: Final = Counter({"D1": 1, "D2": 1, "D3": 1, "D4": 1})
    assert Counter(event.removeprefix("start ") for event in observed if event.startswith("start ")) == configurations
    assert Counter(event.removeprefix("stop ") for event in observed if event.startswith("stop ")) == configurations
    assert sorted(int(event.removeprefix("case ")) for event in observed if event.startswith("case ")) == list(
        range(18)
    )


@pytest.mark.parametrize("workers", (1, 2))
@pytest.mark.parametrize("subset", (False, True))
def test_native_file_selection_reuses_fixtures_and_qualifies_an_empty_shard(
    tmp_path: Path, workers: int, subset: bool
) -> None:
    events: Final = tmp_path / "events.txt"
    root: Final = Path(__file__).resolve().parents[3]
    with TemporaryDirectory(prefix="native-selection-", dir=Path(__file__).parent) as directory:
        probe: Final = Path(directory) / "test_fixture_probe.py"
        plain: Final = Path(directory) / "test_plain_probe.py"
        probe.write_text(
            "import pytest\nfrom pathlib import Path\n"
            f"EVENTS = Path({str(events)!r})\n"
            "def record(value):\n"
            "    with EVENTS.open('a') as output:\n"
            "        output.write(value + '\\n')\n"
            "@pytest.fixture(scope='module')\n"
            "def shared():\n"
            "    record('start shared')\n"
            "    yield 'shared'\n"
            "    record('stop shared')\n"
            "@pytest.mark.parametrize('number', range(4))\n"
            "@pytest.mark.xdist_group('shared')\n"
            "def test_shared(shared, number):\n"
            "    assert shared == 'shared'\n"
            "    record('shared ' + str(number))\n"
        )
        plain.write_text(
            "import pytest\nfrom pathlib import Path\n"
            f"EVENTS = Path({str(events)!r})\n"
            "@pytest.mark.parametrize('number', range(8))\n"
            "def test_plain(number):\n"
            "    with EVENTS.open('a') as output:\n"
            "        output.write('plain ' + str(number) + '\\n')\n"
            "def test_existing_skip():\n"
            "    pytest.skip('retained skip identity')\n"
        )
        selections: Final = (str(probe.relative_to(root)), str(plain.relative_to(root)), "")
        candidates: Final = (
            (f"{selections[0]}::test_shared[1]", f"{selections[1]}::test_existing_skip") if subset else selections[:2]
        )
        listed: Final = subprocess.run(
            (sys.executable, "-P", "tests/integration/run.py", "database", "--list", *candidates),
            cwd=root,
            capture_output=True,
            text=True,
            timeout=15,
            check=False,
        )
        assert listed.returncode == 0, listed.stdout + listed.stderr
        assert tuple(listed.stdout.splitlines()) == tuple(sorted(selections[:2]))
        outputs: Final = tuple(tmp_path / str(index) for index in range(3))
        files: Final = tuple(tmp_path / f"selection-{index}.txt" for index in range(3))
        for selection, path in zip(selections, files, strict=True):
            path.write_text(selection + "\n" if selection else "")
        results: Final = tuple(
            subprocess.run(
                (
                    sys.executable,
                    "-P",
                    "-m",
                    "pytest",
                    *candidates,
                    "--noconftest",
                    "-p",
                    "tests.integration.conftest",
                    "-o",
                    "addopts=",
                    "--integration-shard-count=3",
                    f"--integration-shard-index={index}",
                    f"--integration-shard-files={files[index]}",
                    "--integration-order-seed=1848224111",
                    "--hypothesis-seed=4106601",
                    "--timeout=15",
                    *(("-n", str(workers), "--dist=loadgroup") if workers > 1 else ()),
                    "-q",
                ),
                cwd=root,
                env={
                    **{
                        name: value
                        for name, value in os.environ.items()
                        if name
                        not in (
                            "INTEGRATION_RESULTS_DIR",
                            "INTEGRATION_ROUTING",
                            "GITHUB_ACTIONS",
                            "PYTEST_XDIST_WORKER",
                        )
                    },
                    "INTEGRATION_RESULTS_DIR": str(outputs[index]),
                    "PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), str(root / "tests/e2e"))),
                },
                capture_output=True,
                text=True,
                timeout=30,
                check=False,
            )
            for index in range(3)
        )
        assert all(result.returncode == 0 for result in results), tuple(
            result.stdout + result.stderr for result in results
        )
        executions: Final = tuple(
            ShardExecution.model_validate_json((output / "execution.json").read_text()) for output in outputs
        )
        assert shard_errors(executions, 3) == ()
        assert len(executions[0].inventory) == (2 if subset else 13)
        assert executions[2].collected == ()
        assert executions[2].complete
        assert len(executions[1].skipped) == 1
        assert executions[1].skipped[0].endswith("::test_existing_skip")
    observed: Final = tuple(events.read_text().splitlines())
    assert observed.count("start shared") == observed.count("stop shared") == 1
    assert sorted(event for event in observed if event.startswith("shared ")) == (
        ["shared 1"] if subset else [f"shared {index}" for index in range(4)]
    )
    assert sorted(event for event in observed if event.startswith("plain ")) == (
        [] if subset else [f"plain {index}" for index in range(8)]
    )


@pytest.mark.parametrize("assignment", ("absent", "missing", "foreign"))
def test_invalid_native_selections_fail_before_any_case_runs(
    tmp_path: Path, assignment: Literal["absent", "missing", "foreign"]
) -> None:
    root: Final = Path(__file__).resolve().parents[3]
    events: Final = tmp_path / "executed.txt"
    selection: Final = tmp_path / "node-files.txt"
    if assignment == "foreign":
        selection.write_text("tests/integration/database/test_unknown.py\n")
    with TemporaryDirectory(prefix="invalid-selection-", dir=Path(__file__).parent) as directory:
        probe: Final = Path(directory) / "test_probe.py"
        probe.write_text(f"from pathlib import Path\ndef test_case():\n    Path({str(events)!r}).write_text('ran')\n")
        result: Final = subprocess.run(
            (
                sys.executable,
                "-P",
                "-m",
                "pytest",
                str(probe.relative_to(root)),
                "--noconftest",
                "-p",
                "tests.integration.conftest",
                "-o",
                "addopts=",
                "--integration-shard-count=2",
                "--integration-shard-index=0",
                *((f"--integration-shard-files={selection}",) if assignment != "absent" else ()),
                "-q",
            ),
            cwd=root,
            env={
                **{
                    name: value
                    for name, value in os.environ.items()
                    if name not in ("INTEGRATION_RESULTS_DIR", "INTEGRATION_ROUTING", "GITHUB_ACTIONS")
                },
                "PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), str(root / "tests/e2e"))),
            },
            capture_output=True,
            text=True,
            timeout=15,
            check=False,
        )
        assert result.returncode == pytest.ExitCode.USAGE_ERROR, result.stdout + result.stderr
        assert not events.exists()


@pytest.mark.parametrize("entrypoint", ("gateway", "refused"))
def test_owned_entrypoints_record_exit_and_reap_the_process_group(
    gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, entrypoint: Literal["gateway", "refused"]
) -> None:
    monkeypatch.setenv("INTEGRATION_RESULTS_DIR", str(tmp_path))
    config: Final = tmp_path / "config.yaml"
    if entrypoint == "refused":
        config.write_text("model_list: [\n")
        refused: Final = process_support.refused_boot_log(gateway, tmp_path, {}, config=config)
        assert refused
    else:
        rejected_config: Final = tmp_path / "rejected.yaml"
        rejected_config.write_text("model_list: [\n")
        assert process_support.refused_boot_log(gateway, tmp_path, {}, config=rejected_config)
        config.write_text(
            "model_list: []\ngeneral_settings:\n  master_key: os.environ/LITELLM_MASTER_KEY\n"
            "  database_url: os.environ/DATABASE_URL\n"
        )
        with process_support.owned_gateway_image(gateway, tmp_path, {}, config=config, workers=1) as owned:
            response: Final = owned.gateway.request("GET", "/health/readiness")
            assert response.status_code == 200, response.text
        completed: Final = TypeAdapter(ExitRecord).validate_json(owned.log.with_suffix(".exit.json").read_text())
        assert completed["root_pid"] == owned.process.pid
        assert completed["exit_after_stop"] == owned.process.returncode
    records: Final = tuple(
        TypeAdapter(ExitRecord).validate_json(path.read_text()) for path in tmp_path.glob("*.exit.json")
    )
    assert records
    if entrypoint == "refused":
        assert len(records) == 1 and records[0]["exit_after_stop"] != 0
    for record in records:
        assert record["exit_after_stop"] is not None
        assert record["exception_at_record"] is None
        assert not process_support.group_members(record["root_pid"])