next-tjwater-server/tests/api/test_health_endpoints.py

import asyncio
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

from app.api import health


@pytest.fixture
def probe_app(monkeypatch):
    app = FastAPI()
    app.include_router(health.router)
    app.state.startup_complete = True
    app.state.mcp_runtime = SimpleNamespace(ready=True)
    database_check = AsyncMock(return_value=True)
    monkeypatch.setattr(health, "metadata_database_ready", database_check)
    return app, TestClient(app), database_check


@pytest.mark.parametrize("runtime", [None, SimpleNamespace(ready=True)])
def test_ready_accepts_disabled_or_started_mcp(probe_app, runtime):
    app, client, database_check = probe_app
    app.state.mcp_runtime = runtime
    response = client.get("/ready")
    assert response.status_code == 200
    assert response.json() == {
        "status": "ready",
        "checks": {
            "startup": "ready", "metadata_database": "ready",
            "mcp_transport": "disabled" if runtime is None else "ready",
        },
    }
    assert response.headers["cache-control"] == "no-store"
    database_check.assert_awaited_once_with()


def test_startup_and_shutdown_are_not_ready_without_touching_database(probe_app):
    app, client, database_check = probe_app
    app.state.startup_complete = False
    response = client.get("/ready")
    assert response.status_code == 503
    assert response.json()["checks"]["startup"] == "not_ready"
    assert response.json()["checks"]["metadata_database"] == "not_checked"
    database_check.assert_not_awaited()


@pytest.mark.parametrize("component", ["metadata_database", "mcp_transport"])
def test_dependency_not_ready_returns_503_without_breaking_liveness(probe_app, component):
    app, client, database_check = probe_app
    if component == "metadata_database":
        database_check.return_value = False
    else:
        app.state.mcp_runtime.ready = False
    response = client.get("/ready")
    assert response.status_code == 503
    assert response.json()["status"] == "not_ready"
    assert response.json()["checks"][component] in {"unavailable", "not_ready"}
    database_check.reset_mock()
    response = client.get("/health")
    assert response.status_code == 200
    assert response.json() == {"status": "ok"}
    database_check.assert_not_awaited()


@pytest.mark.parametrize("outcome", ["ready", "error", "timeout"])
def test_metadata_probe_runs_bounded_select_and_sanitizes_errors(monkeypatch, outcome):
    queries = []
    closed = []

    async def execute(query):
        queries.append(str(query))
        if outcome == "error":
            raise RuntimeError("postgresql://private:secret@internal/database")
        if outcome == "timeout":
            await asyncio.Event().wait()
        return SimpleNamespace(scalar_one=lambda: 1)

    @asynccontextmanager
    async def connect():
        try:
            yield SimpleNamespace(execute=execute)
        finally:
            closed.append(True)

    monkeypatch.setattr(health, "engine", SimpleNamespace(connect=connect))
    monkeypatch.setattr(health, "METADATA_TIMEOUT_SECONDS", 0.01)
    assert asyncio.run(health.metadata_database_ready()) is (outcome == "ready")
    assert queries == ["SELECT 1"]
    assert closed == [True]


@pytest.mark.parametrize("startup_fails", [False, True])
def test_application_lifecycle_marks_readiness_before_cleanup(monkeypatch, startup_fails):
    from app import main

    app = FastAPI()

    @asynccontextmanager
    async def mcp_lifespan():
        assert app.state.startup_complete is False
        if startup_fails:
            raise RuntimeError("MCP startup failed")
        try:
            yield
        finally:
            assert app.state.startup_complete is False

    monkeypatch.setattr(main, "mcp_runtime", SimpleNamespace(lifespan=mcp_lifespan))
    for owner, method in [
        (main.workbench_event_hubs, "close_all"),
        (main.project_connection_manager, "close_all"),
        (main, "close_metadata_engine"),
    ]:
        monkeypatch.setattr(owner, method, AsyncMock())
    monkeypatch.setattr(main, "close_all_timescale_pools", Mock())
    monkeypatch.setattr(main, "close_all_project_pools", Mock())

    async def run():
        if startup_fails:
            with pytest.raises(RuntimeError, match="MCP startup failed"):
                async with main.lifespan(app):
                    pytest.fail("Failed startup must not yield")
        else:
            async with main.lifespan(app):
                assert app.state.startup_complete is True
        assert app.state.startup_complete is False

    asyncio.run(run())


def test_mcp_runtime_tracks_its_actual_lifecycle():
    from app.mcp.application import McpQueryApplication
    from app.mcp.config import McpSettings

    runtime = McpQueryApplication(McpSettings(
        enabled=True, issuer_uri="http://localhost/realms/water",
        resource_uri="http://localhost/mcp", _env_file=None,
    ))

    async def run():
        assert runtime.ready is False
        async with runtime.lifespan():
            assert runtime.ready is True
        assert runtime.ready is False

    asyncio.run(run())