# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""
Tests for DPSupervisor: unit tests and lifecycle integration tests.

Lifecycle integration tests replace child vLLM servers with lightweight
aiohttp "fake" servers controlled by the test, so the suite runs without GPUs.
_start_children is monkeypatched to install FakeProcess objects (with
controllable liveness/timing) alongside those fake HTTP servers.

Port allocation (kept far from default vLLM ports to avoid conflicts):
  Supervisor : 19256
  Children   : 18000, 18001
"""

from __future__ import annotations

import argparse
import asyncio
import contextlib
import os
import signal
import subprocess
import tempfile
import time
from pathlib import Path
from types import SimpleNamespace

import aiohttp
import pytest
import uvicorn
from fastapi import FastAPI, Response

import vllm.entrypoints.openai.dp_supervisor as dp_sup
from vllm.entrypoints.openai.dp_supervisor import (
    CHILD_EXIT_GRACE_S,
    DPSupervisor,
    _build_vllm_dp_server_args,
    infer_multi_port_external_lb_start_rank,
    validate_multi_port_external_lb_args,
)
from vllm.logger import init_logger

logger = init_logger(__name__)

_SUPERVISOR_PORT = 19256
_CHILD_PORT_BASE = 18000
_N_CHILDREN = 2
_PROBE_INTERVAL = 1.0
_POLL_INTERVAL = 1.0


# ---------------------------------------------------------------------------
# Args factories
# ---------------------------------------------------------------------------


def _make_unit_args(**overrides) -> argparse.Namespace:
    """Minimal args for unit tests (no real network activity)."""
    base = {
        "host": None,
        "port": 8000,
        "data_parallel_multi_port_external_lb": True,
        "data_parallel_supervisor_port": 9256,
        "dp_supervisor_probe_interval_s": 5.0,
        "dp_supervisor_probe_timeout_s": 5.0,
        "dp_supervisor_probe_failure_threshold": 3,
        "data_parallel_size": 8,
        "data_parallel_size_local": 4,
        "data_parallel_start_rank": None,
        "data_parallel_rank": None,
        "data_parallel_external_lb": False,
        "data_parallel_hybrid_lb": False,
        "api_server_count": None,
        "headless": False,
        "grpc": False,
        "uds": None,
        "ssl_keyfile": None,
        "ssl_certfile": None,
        "ssl_ca_certs": None,
        "ssl_cert_reqs": 0,
        "ssl_ciphers": None,
        "node_rank": 1,
        "tensor_parallel_size": 1,
        "pipeline_parallel_size": 1,
        "uvicorn_log_level": "info",
        "shutdown_timeout": 5.0,
    }
    base.update(overrides)
    return argparse.Namespace(**base)


def _make_args(**overrides) -> argparse.Namespace:
    """Args for lifecycle integration tests (real loopback servers)."""
    base: dict = dict(
        host="127.0.0.1",
        port=_CHILD_PORT_BASE,
        data_parallel_multi_port_external_lb=True,
        data_parallel_supervisor_port=_SUPERVISOR_PORT,
        dp_supervisor_probe_interval_s=_PROBE_INTERVAL,
        dp_supervisor_probe_timeout_s=1.0,
        dp_supervisor_probe_failure_threshold=3,
        data_parallel_size=_N_CHILDREN,
        data_parallel_size_local=_N_CHILDREN,
        data_parallel_start_rank=0,
        data_parallel_rank=None,
        data_parallel_external_lb=False,
        data_parallel_hybrid_lb=False,
        api_server_count=None,
        headless=False,
        grpc=False,
        uds=None,
        ssl_keyfile=None,
        ssl_certfile=None,
        ssl_ca_certs=None,
        ssl_cert_reqs=0,
        ssl_ciphers=None,
        node_rank=0,
        tensor_parallel_size=1,
        pipeline_parallel_size=1,
        uvicorn_log_level="warning",
        shutdown_timeout=0.0,
    )
    base.update(overrides)
    return argparse.Namespace(**base)


def _generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
    """Generate a self-signed certificate for HTTPS lifecycle tests."""
    cert_file = cert_dir / "cert.pem"
    key_file = cert_dir / "key.pem"
    subprocess.run(
        [
            "openssl",
            "req",
            "-x509",
            "-newkey",
            "rsa:2048",
            "-keyout",
            str(key_file),
            "-out",
            str(cert_file),
            "-days",
            "1",
            "-nodes",
            "-subj",
            "/CN=localhost",
        ],
        check=True,
        capture_output=True,
    )
    return cert_file, key_file


# ---------------------------------------------------------------------------
# Unit tests
# ---------------------------------------------------------------------------


def test_infer_multi_port_external_lb_start_rank_uses_node_rank():
    args = _make_unit_args()
    assert infer_multi_port_external_lb_start_rank(args) == 4


def test_build_multi_port_external_lb_child_args_sets_external_rank_server():
    args = _make_unit_args(data_parallel_start_rank=8, api_server_count=None)
    child_args = _build_vllm_dp_server_args(args, local_rank=2)

    assert child_args.port == 8002
    assert child_args.data_parallel_rank == 10
    assert child_args.data_parallel_size_local == 1
    assert child_args.data_parallel_external_lb is True
    assert child_args.data_parallel_hybrid_lb is False
    assert child_args.data_parallel_multi_port_external_lb is False
    assert child_args.api_server_count == 1


def test_run_vllm_dp_server_uses_python_server_by_default(monkeypatch):
    calls: list[str] = []

    monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
    monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
    monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
    monkeypatch.setattr(dp_sup.envs, "VLLM_RUST_FRONTEND_PATH", None, raising=False)
    monkeypatch.setattr(
        dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
    )
    monkeypatch.setattr(
        dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
    )

    dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))

    assert calls == ["python"]


def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch):
    calls: list[str] = []

    monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
    monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
    monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
    monkeypatch.setattr(
        dp_sup.envs,
        "VLLM_RUST_FRONTEND_PATH",
        "/tmp/vllm-rs",
        raising=False,
    )
    monkeypatch.setattr(
        dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
    )
    monkeypatch.setattr(
        dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
    )

    dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))

    assert calls == ["rust"]


def test_validate_multi_port_external_lb_args_allows_ssl():
    args = _make_unit_args(
        ssl_keyfile="/tmp/server.key",
        ssl_certfile="/tmp/server.crt",
        ssl_ca_certs="/tmp/ca.crt",
    )
    validate_multi_port_external_lb_args(args)


def test_aggregates_health():
    supervisor = DPSupervisor(_make_unit_args())
    supervisor._is_ready = True
    assert supervisor.is_ready is True


def test_handles_shutdown_event():
    supervisor = DPSupervisor(_make_unit_args())
    supervisor._is_ready = True
    supervisor._shutdown_event.set()
    assert supervisor.is_ready is False


@pytest.mark.asyncio
async def test_handles_child_exit(
    monkeypatch: pytest.MonkeyPatch,
):
    supervisor = DPSupervisor(_make_unit_args())
    supervisor._processes = [
        SimpleNamespace(
            name="APIServer_DPRank_4", exitcode=None, is_alive=lambda: True
        ),
        SimpleNamespace(name="APIServer_DPRank_5", exitcode=17, is_alive=lambda: False),
    ]

    async def fake_probe(*_args, **_kwargs) -> bool:
        return True

    monkeypatch.setattr(dp_sup, "_probe_endpoint", fake_probe)

    await supervisor._monitor_children()
    assert supervisor._is_ready is False


@pytest.mark.asyncio
async def test_handles_probe_failure(
    monkeypatch: pytest.MonkeyPatch,
):
    supervisor = DPSupervisor(_make_unit_args(dp_supervisor_probe_interval_s=0.0))
    supervisor.child_ports = [8000]
    probe_results = iter([True, False])

    async def fake_probe(*_args, **_kwargs) -> bool:
        return next(probe_results)

    monkeypatch.setattr(dp_sup, "_probe_endpoint", fake_probe)

    await supervisor._monitor_children()
    assert supervisor._is_ready is False


@pytest.mark.asyncio
async def test_shutdown_if_supervisor_server_error_on_startup(
    monkeypatch: pytest.MonkeyPatch,
):
    class FakeLoop:
        def add_signal_handler(self, *_args, **_kwargs):
            pass

        def remove_signal_handler(self, *_args, **_kwargs):
            pass

    class FakeServer:
        def __init__(self, _config):
            self.started = False
            self.should_exit = False

        async def serve(self):
            raise ValueError("supervisor boom")

    async def fake_shutdown_children(self):
        return None

    def fake_start_children(self):
        return None

    async def fake_monitor_children(self):
        # Mark ready so the supervisor server is started, then block until
        # shutdown (triggered when the failing server task exits).
        self._is_ready = True
        await self._shutdown_event.wait()

    monkeypatch.setattr(dp_sup.asyncio, "get_running_loop", lambda: FakeLoop())
    monkeypatch.setattr(dp_sup.uvicorn, "Server", FakeServer)
    monkeypatch.setattr(DPSupervisor, "_shutdown_children", fake_shutdown_children)
    monkeypatch.setattr(DPSupervisor, "_start_children", fake_start_children)
    monkeypatch.setattr(DPSupervisor, "_monitor_children", fake_monitor_children)

    supervisor = DPSupervisor(_make_unit_args())

    with pytest.raises(ValueError, match="supervisor boom"):
        await supervisor.run()


# ---------------------------------------------------------------------------
# Lifecycle integration tests – MockVLLMServer
# ---------------------------------------------------------------------------


class MockVLLMServer:
    """
    Minimal FastAPI server that mimics one vLLM replica.
    GET /health returns 200 when healthy, 503 otherwise.
    Health state is toggled by the test via set_healthy().
    """

    def __init__(
        self,
        port: int,
        drain_seconds: float = 0.0,
        ssl_keyfile: str | None = None,
        ssl_certfile: str | None = None,
    ) -> None:
        self.port = port
        self._healthy = False
        self._drain_seconds = drain_seconds
        self._ssl_keyfile = ssl_keyfile
        self._ssl_certfile = ssl_certfile
        self._server: uvicorn.Server | None = None
        self._serve_task: asyncio.Task | None = None

    async def start(self) -> None:
        app = FastAPI()

        @app.get("/health")
        async def health() -> Response:
            print(f"MockServer {self.port}: /health: {self._healthy}")
            return Response(status_code=200 if self._healthy else 503)

        @app.get("/set_healthy")
        async def set_healthy() -> Response:
            print(f"MockServer {self.port}: /set_healthy")
            self._healthy = True
            return Response(status_code=200)

        @app.get("/set_unhealthy")
        async def set_unhealthy() -> Response:
            print(f"MockServer {self.port}: /set_unhealthy")
            self._healthy = False
            return Response(status_code=200)

        @app.get("/kill")
        async def kill() -> Response:
            print(f"MockServer {self.port}: /kill")
            os.kill(os.getpid(), signal.SIGKILL)

        config = uvicorn.Config(
            app,
            host="127.0.0.1",
            port=self.port,
            log_level="warning",
            lifespan="off",
            ssl_keyfile=self._ssl_keyfile,
            ssl_certfile=self._ssl_certfile,
        )
        self._server = uvicorn.Server(config)

        # Configure request draining if needed.
        # Uvicorn's capture_signals() installs signal.signal(SIGTERM, self.handle_exit),
        # which sets should_exit=True immediately. Override handle_exit on the instance
        # so capture_signals() picks up our version that drains first.
        if self._drain_seconds > 0:
            self._shutdown_event = asyncio.Event()
            loop = asyncio.get_running_loop()

            async def _drain_and_stop() -> None:
                await self._shutdown_event.wait()
                print(f"MockServer {self.port}: draining for {self._drain_seconds}s.")
                await asyncio.sleep(self._drain_seconds)
                print("Setting should_exit")
                if self._server is not None:
                    self._server.should_exit = True

            self._drain_task = asyncio.create_task(_drain_and_stop())

            def _custom_handle_exit(sig: int, frame: object) -> None:
                print("Got SIGTERM, setting shutdown.")
                if not self._shutdown_event.is_set():
                    loop.call_soon_threadsafe(self._shutdown_event.set)

            self._server.handle_exit = _custom_handle_exit

        self._serve_task = asyncio.create_task(self._server.serve())
        while not self._server.started:
            await asyncio.sleep(0.01)
        print(f"Mock DP Server on port {self.port} started")

        await self._serve_task


def launch_mock_vllm(child_args: argparse.Namespace):
    logger.info("Launching mock vLLM on port %s", child_args.port)
    mock_vllm = MockVLLMServer(
        port=child_args.port,
        ssl_keyfile=child_args.ssl_keyfile,
        ssl_certfile=child_args.ssl_certfile,
    )
    asyncio.run(mock_vllm.start())


def launch_mock_vllm_with_drain(
    child_args: argparse.Namespace,
):
    logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
    mock_vllm = MockVLLMServer(
        port=child_args.port,
        drain_seconds=10.0,
        ssl_keyfile=child_args.ssl_keyfile,
        ssl_certfile=child_args.ssl_certfile,
    )
    asyncio.run(mock_vllm.start())


# ---------------------------------------------------------------------------
# Lifecycle test helpers
# ---------------------------------------------------------------------------


async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool:
    """
    GET /health on the supervisor once and check for expected_status.

    Pass expected_status=-1 to assert the supervisor is not listening yet
    (a connection error is expected). The supervisor only starts its HTTP
    server once every child is ready, so /health is refused until then.
    """
    scheme = "https" if use_ssl else "http"
    url = f"{scheme}://127.0.0.1:{_SUPERVISOR_PORT}/health"
    async with aiohttp.ClientSession() as session:
        try:
            async with session.get(url, ssl=False if use_ssl else None) as resp:
                if resp.status != expected_status:
                    print(f"expected: {expected_status=}, got: {resp.status=}")
                    return False
                return True
        except aiohttp.ClientError:
            if expected_status != -1:
                print(f"expected: {expected_status=}, got: aiohttp.ClientError")
                return False
            return True


async def _await_supervisor_health(
    expected_status: int, use_ssl: bool = False, retries: int = 20
) -> bool:
    """Retry _poll_supervisor_health, tolerating supervisor server startup."""
    for _ in range(retries):
        if await _poll_supervisor_health(expected_status, use_ssl=use_ssl):
            return True
        await asyncio.sleep(0.5)
    return False


async def _poll_until_api_server_running(
    port: int, retries: int = 10, use_ssl: bool = False
) -> None:
    scheme = "https" if use_ssl else "http"
    url = f"{scheme}://127.0.0.1:{port}/health"
    async with aiohttp.ClientSession() as session:
        for _ in range(retries):
            try:
                async with session.get(url, ssl=False if use_ssl else None) as resp:
                    if resp.status != 200:
                        return
                await asyncio.sleep(1.0)
            except aiohttp.ClientError:
                print("Test detected not started yet, sleeping for 1s")
                await asyncio.sleep(1.0)


async def _set_healthy(port: int, use_ssl: bool = False) -> None:
    scheme = "https" if use_ssl else "http"
    url = f"{scheme}://127.0.0.1:{port}/set_healthy"
    async with (
        aiohttp.ClientSession() as session,
        session.get(url, ssl=False if use_ssl else None) as resp,
    ):
        assert resp.status == 200


async def _set_unhealthy(port: int, use_ssl: bool = False) -> None:
    scheme = "https" if use_ssl else "http"
    url = f"{scheme}://127.0.0.1:{port}/set_unhealthy"
    async with (
        aiohttp.ClientSession() as session,
        session.get(url, ssl=False if use_ssl else None) as resp,
    ):
        assert resp.status == 200


async def _kill_server(port: int, use_ssl: bool = False) -> None:
    scheme = "https" if use_ssl else "http"
    url = f"{scheme}://127.0.0.1:{port}/kill"
    try:
        async with (
            aiohttp.ClientSession() as session,
            session.get(url, ssl=False if use_ssl else None) as resp,
        ):
            assert resp.status != 200
    except Exception as e:
        assert isinstance(e, aiohttp.ClientConnectorError)


@contextlib.asynccontextmanager
async def _run_supervisor(
    args: argparse.Namespace,
    monkeypatch: pytest.MonkeyPatch,
    launch_fn=None,
):
    if launch_fn is None:
        launch_fn = launch_mock_vllm
    monkeypatch.setattr(dp_sup, "_run_vllm_dp_server", launch_fn)
    supervisor = DPSupervisor(args)
    task = asyncio.create_task(supervisor.run())
    await asyncio.sleep(1.0)
    try:
        yield supervisor, task
    finally:
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            await task


# ---------------------------------------------------------------------------
# Lifecycle integration tests
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_basic_lifecycle(monkeypatch):
    """
    A) Supervisor is not listening while children are unhealthy.
    B) /health returns 200 once every child reports healthy.
    C) SIGTERM and shutdown
    """
    args = _make_args()

    vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]

    async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
        assert await _poll_supervisor_health(-1)
        assert not supervisor.is_ready

        for port in vllm_server_ports:
            assert await _poll_supervisor_health(-1)
            assert not supervisor.is_ready
            await _poll_until_api_server_running(port)

        await _set_healthy(vllm_server_ports[0])
        await asyncio.sleep(1.0)
        assert await _poll_supervisor_health(-1)
        assert not supervisor.is_ready
        print("supervisor not listening --- expected!")

        for port in vllm_server_ports:
            await _set_healthy(port)
        assert await _await_supervisor_health(200)
        assert supervisor.is_ready
        print("/health is 200 --- expected!")

        await asyncio.sleep(1.0)
        assert await _poll_supervisor_health(200)
        assert supervisor.is_ready
        print("/health is 200 --- expected!")

        os.kill(os.getpid(), signal.SIGTERM)

        await asyncio.wait_for(_task, timeout=5.0)
        for p in supervisor._processes:
            assert not p.is_alive()
        print("everything was cleaned up!")


@pytest.mark.asyncio
async def test_basic_lifecycle_with_ssl(monkeypatch):
    with tempfile.TemporaryDirectory() as cert_dir:
        cert_file, key_file = _generate_self_signed_cert(Path(cert_dir))
        args = _make_args(
            ssl_keyfile=str(key_file),
            ssl_certfile=str(cert_file),
        )

        vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]

        async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
            assert await _poll_supervisor_health(-1, use_ssl=True)
            assert not supervisor.is_ready

            for port in vllm_server_ports:
                assert await _poll_supervisor_health(-1, use_ssl=True)
                assert not supervisor.is_ready
                await _poll_until_api_server_running(port, use_ssl=True)

            for port in vllm_server_ports:
                await _set_healthy(port, use_ssl=True)

            assert await _await_supervisor_health(200, use_ssl=True)
            assert supervisor.is_ready


@pytest.mark.asyncio
async def test_failed_startup(monkeypatch):
    """
    A) One of the vLLM servers crashes during startup.
    B) DPSupervisor detects this, and cleans up resources.
    """
    args = _make_args()

    vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]

    async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
        assert await _poll_supervisor_health(-1)
        assert not supervisor.is_ready

        for port in vllm_server_ports:
            await _poll_until_api_server_running(port)

        await _kill_server(port)

        await asyncio.wait_for(_task, timeout=5.0)
        for p in supervisor._processes:
            assert not p.is_alive()


@pytest.mark.asyncio
async def test_becomes_unhealthy(monkeypatch):
    """
    A) Supervisor is not listening while children are unhealthy.
    B) /health returns 200 once every child reports healthy.
    C) Child process becomes unhealthy.
    D) Detected and shutdown.
    """
    args = _make_args()

    vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]

    async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
        assert await _poll_supervisor_health(-1)
        assert not supervisor.is_ready

        for port in vllm_server_ports:
            assert await _poll_supervisor_health(-1)
            assert not supervisor.is_ready
            await _poll_until_api_server_running(port)

        await _set_healthy(vllm_server_ports[0])
        await asyncio.sleep(1.0)
        assert await _poll_supervisor_health(-1)
        assert not supervisor.is_ready
        print("supervisor not listening --- expected!")

        for port in vllm_server_ports:
            await _set_healthy(port)
        assert await _await_supervisor_health(200)
        assert supervisor.is_ready
        print("/health is 200 --- expected!")

        await _set_unhealthy(port)

        await asyncio.wait_for(_task, timeout=5.0)
        for p in supervisor._processes:
            assert not p.is_alive()
        print("everything was cleaned up!")


@pytest.mark.asyncio
async def test_dp_server_fails(monkeypatch):
    """
    A) Supervisor is not listening while children are unhealthy.
    B) /health returns 200 once every child reports healthy.
    C) Child process fails.
    D) Detected and shutdown.
    """
    args = _make_args()

    vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]

    async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
        assert await _poll_supervisor_health(-1)
        assert not supervisor.is_ready

        for port in vllm_server_ports:
            assert await _poll_supervisor_health(-1)
            assert not supervisor.is_ready
            await _poll_until_api_server_running(port)

        await _set_healthy(vllm_server_ports[0])
        await asyncio.sleep(1.0)
        assert await _poll_supervisor_health(-1)
        assert not supervisor.is_ready
        print("supervisor not listening --- expected!")

        for port in vllm_server_ports:
            await _set_healthy(port)
        assert await _await_supervisor_health(200)
        assert supervisor.is_ready
        print("/health is 200 --- expected!")

        dp_mock_server_process = supervisor._processes[0]
        os.kill(dp_mock_server_process.pid, signal.SIGKILL)
        await asyncio.sleep(1.0)
        assert not dp_mock_server_process.is_alive()

        await asyncio.wait_for(_task, timeout=5.0)
        for p in supervisor._processes:
            assert not p.is_alive()
        print("everything was cleaned up!")


@pytest.mark.asyncio
async def test_shutdown_timeout(monkeypatch: pytest.MonkeyPatch):
    """
    Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight
    request drain).  The supervisor is configured with shutdown_timeout=10,
    so its total wait budget is 10 + CHILD_EXIT_GRACE_S seconds.  The
    children exit naturally within that window, so no force-kill should occur
    and the measured wall-clock time must be >= 10s.
    """
    _DRAIN_SECONDS = 10.0
    _SHUTDOWN_TIMEOUT = 10.0

    args = _make_args(shutdown_timeout=_SHUTDOWN_TIMEOUT)
    vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]

    async with _run_supervisor(
        args, monkeypatch, launch_fn=launch_mock_vllm_with_drain
    ) as (supervisor, _task):
        for port in vllm_server_ports:
            await _poll_until_api_server_running(port)

        for port in vllm_server_ports:
            await _set_healthy(port)
        assert await _await_supervisor_health(200)
        assert supervisor.is_ready

        start_t = time.perf_counter()
        os.kill(os.getpid(), signal.SIGTERM)

        print(f"DRAINING FOR {_DRAIN_SECONDS}")
        await asyncio.wait_for(_task, timeout=_DRAIN_SECONDS + CHILD_EXIT_GRACE_S + 5.0)
        elapsed = time.perf_counter() - start_t

        assert elapsed >= _DRAIN_SECONDS, (
            f"Supervisor exited after only {elapsed:.1f}s; "
            f"expected >= {_DRAIN_SECONDS}s for request draining"
        )

        for p in supervisor._processes:
            assert not p.is_alive()
        print(f"Supervisor waited {elapsed:.1f}s for children to drain — expected!")