FabledSteward/tests/core/test_db_readiness.py

"""DB connection resilience: pool, startup readiness, and request-time outage.

Covers the three recovery gaps that let a database restart take the app down:
the engine handing out stale pooled connections (#2626), startup assuming the
DB is reachable and ready the instant it asks (#2627), and a request finding
the database gone and failing with an opaque 500 (#2635).
"""
import types
import pytest

from steward import database
from steward.database import (
    POOL_RECYCLE_SECONDS,
    _describe,
    _fatal_sqlstate,
    init_db,
    wait_for_database,
)


class _PGError(Exception):
    """Stand-in for an asyncpg error, which carries a SQLSTATE attribute."""

    def __init__(self, message: str, sqlstate: str | None = None):
        super().__init__(message)
        self.sqlstate = sqlstate


@pytest.fixture
def no_sleep(monkeypatch):
    """Drive the retry budget off a fake clock instead of wall time.

    sleep() records its duration and advances the clock by exactly that much,
    so backoff and timeout behaviour are deterministic rather than dependent on
    how fast the test machine happens to run.
    """
    slept: list[float] = []
    clock = {"now": 0.0}

    def fake_sleep(seconds: float) -> None:
        slept.append(seconds)
        clock["now"] += seconds

    monkeypatch.setattr(database.time, "sleep", fake_sleep)
    monkeypatch.setattr(database.time, "monotonic", lambda: clock["now"])
    return slept


def _probe_raising(*errors, then_succeed: bool = True):
    """Async probe stub raising the given errors in order.

    Once the list is exhausted it succeeds, unless then_succeed is False — in
    which case it keeps raising the last error forever (a DB that never comes
    back).
    """
    calls = {"n": 0}

    async def probe(db_url):
        i = calls["n"]
        calls["n"] += 1
        if i < len(errors):
            raise errors[i]
        if not then_succeed:
            raise errors[-1]

    probe.calls = calls
    return probe


# ── engine pool configuration (issue #2626) ──────────────────────────────────


def test_engine_enables_pre_ping_and_recycle():
    app = types.SimpleNamespace(
        config={"DATABASE_URL": "postgresql+asyncpg://u:p@localhost/db"}
    )
    init_db(app)

    # _pre_ping / _recycle are SQLAlchemy pool internals; there is no public
    # accessor, and these are exactly the settings a DB restart depends on.
    pool = app._db_engine.sync_engine.pool
    assert pool._pre_ping is True
    assert pool._recycle == POOL_RECYCLE_SECONDS


def test_pool_recycle_is_positive():
    # A non-positive recycle disables age-based retirement entirely.
    assert POOL_RECYCLE_SECONDS > 0


# ── retry classification ─────────────────────────────────────────────────────


def test_fatal_sqlstate_detects_bad_password():
    exc = _PGError("password authentication failed", sqlstate="28P01")
    assert _fatal_sqlstate(exc) == "28P01"


def test_fatal_sqlstate_finds_sqlstate_through_cause_chain():
    inner = _PGError("database does not exist", sqlstate="3D000")
    outer = RuntimeError("wrapped by sqlalchemy")
    outer.__cause__ = inner
    assert _fatal_sqlstate(outer) == "3D000"


def test_transient_errors_are_not_fatal():
    assert _fatal_sqlstate(OSError(-2, "Name or service not known")) is None
    # "the database system is not yet accepting connections" — resolves on its own.
    assert _fatal_sqlstate(_PGError("not yet accepting", sqlstate="57P03")) is None


def test_describe_unwraps_to_innermost_cause():
    inner = OSError("Name or service not known")
    outer = RuntimeError("sqlalchemy boilerplate")
    outer.__cause__ = inner
    assert "Name or service not known" in _describe(outer)


def test_describe_survives_self_referential_cause():
    exc = RuntimeError("loop")
    exc.__cause__ = exc
    assert "loop" in _describe(exc)


# ── wait_for_database (issue #2627) ──────────────────────────────────────────


def test_returns_immediately_when_db_is_up(monkeypatch, no_sleep):
    probe = _probe_raising()
    monkeypatch.setattr(database, "_probe", probe)

    wait_for_database("postgresql+asyncpg://u:p@db/steward", timeout_seconds=10)

    assert probe.calls["n"] == 1
    assert no_sleep == []


def test_retries_dns_failure_then_succeeds(monkeypatch, no_sleep):
    # The reported failure: gaierror -2 while Docker DNS is not yet serving `db`.
    probe = _probe_raising(
        OSError(-2, "Name or service not known"),
        OSError(-2, "Name or service not known"),
    )
    monkeypatch.setattr(database, "_probe", probe)

    wait_for_database("postgresql+asyncpg://u:p@db/steward", timeout_seconds=30)

    assert probe.calls["n"] == 3
    assert len(no_sleep) == 2


def test_retries_while_postgres_is_still_recovering(monkeypatch, no_sleep):
    probe = _probe_raising(
        _PGError("the database system is not yet accepting connections", "57P03"),
    )
    monkeypatch.setattr(database, "_probe", probe)

    wait_for_database("postgresql+asyncpg://u:p@db/steward", timeout_seconds=30)

    assert probe.calls["n"] == 2


def test_backoff_grows_between_attempts(monkeypatch, no_sleep):
    probe = _probe_raising(*[OSError("refused")] * 4)
    monkeypatch.setattr(database, "_probe", probe)

    wait_for_database("postgresql+asyncpg://u:p@db/steward", timeout_seconds=60)

    assert no_sleep == sorted(no_sleep), "delays should be non-decreasing"
    assert no_sleep[-1] > no_sleep[0], "backoff should grow, not stay flat"


def test_backoff_is_capped(monkeypatch, no_sleep):
    probe = _probe_raising(*[OSError("refused")] * 12)
    monkeypatch.setattr(database, "_probe", probe)

    wait_for_database("postgresql+asyncpg://u:p@db/steward", timeout_seconds=600)

    assert max(no_sleep) <= database._BACKOFF_MAX_SECONDS


def test_bad_credentials_fail_immediately_without_retrying(monkeypatch, no_sleep):
    probe = _probe_raising(_PGError("password authentication failed", "28P01"))
    monkeypatch.setattr(database, "_probe", probe)

    with pytest.raises(_PGError):
        wait_for_database("postgresql+asyncpg://u:bad@db/steward", timeout_seconds=30)

    # Waiting cannot fix a wrong password — don't burn the whole budget on it.
    assert probe.calls["n"] == 1
    assert no_sleep == []


def test_missing_database_fails_immediately(monkeypatch, no_sleep):
    probe = _probe_raising(_PGError("database does not exist", "3D000"))
    monkeypatch.setattr(database, "_probe", probe)

    with pytest.raises(_PGError):
        wait_for_database("postgresql+asyncpg://u:p@db/nope", timeout_seconds=30)

    assert probe.calls["n"] == 1


def test_gives_up_after_budget_with_actionable_message(monkeypatch, no_sleep):
    probe = _probe_raising(
        OSError("Name or service not known"), then_succeed=False
    )
    monkeypatch.setattr(database, "_probe", probe)

    with pytest.raises(RuntimeError) as excinfo:
        wait_for_database("postgresql+asyncpg://u:p@db/steward", timeout_seconds=2)

    message = str(excinfo.value)
    assert "did not become available" in message
    # The operator needs the underlying cause, not just "timed out".
    assert "Name or service not known" in message


# ── ensure_database_reachable (issue #2635) ──────────────────────────────────


class _FakeEngine:
    """Minimal stand-in for AsyncEngine.connect() as an async context manager."""

    def __init__(self, *errors, then_succeed: bool = True):
        self.errors = list(errors)
        self.then_succeed = then_succeed
        self.calls = 0

    def connect(self):
        engine = self

        class _Ctx:
            async def __aenter__(self):
                index = engine.calls
                engine.calls += 1
                if index < len(engine.errors):
                    raise engine.errors[index]
                if not engine.then_succeed:
                    raise engine.errors[-1]
                return object()

            async def __aexit__(self, *exc_info):
                return False

        return _Ctx()


@pytest.fixture
def no_async_sleep(monkeypatch):
    """Record awaited retry delays without spending them."""
    slept: list[float] = []

    async def fake_sleep(seconds: float) -> None:
        slept.append(seconds)

    monkeypatch.setattr(database.asyncio, "sleep", fake_sleep)
    return slept


async def test_reachable_on_first_try_costs_no_retries(no_async_sleep):
    engine = _FakeEngine()
    await database.ensure_database_reachable(engine)
    assert engine.calls == 1
    assert no_async_sleep == []


async def test_rides_out_a_brief_outage(no_async_sleep):
    # The reported shape: DNS gone while the db container restarts.
    engine = _FakeEngine(OSError(-2, "Name or service not known"))
    await database.ensure_database_reachable(engine)
    assert engine.calls == 2
    assert no_async_sleep == [0.25]


async def test_uses_the_full_retry_schedule_before_giving_up(no_async_sleep):
    engine = _FakeEngine(OSError("refused"), then_succeed=False)
    with pytest.raises(database.DatabaseUnavailable):
        await database.ensure_database_reachable(engine)

    # One attempt per delay, plus a final attempt after the last wait.
    assert engine.calls == len(database._REQUEST_RETRY_DELAYS) + 1
    assert no_async_sleep == list(database._REQUEST_RETRY_DELAYS)


async def test_request_budget_stays_short(no_async_sleep):
    # A person is waiting on this; guard against the schedule growing into a hang.
    assert sum(database._REQUEST_RETRY_DELAYS) <= 3.0


async def test_unavailable_error_carries_the_underlying_cause(no_async_sleep):
    engine = _FakeEngine(
        OSError("Name or service not known"), then_succeed=False)
    with pytest.raises(database.DatabaseUnavailable) as excinfo:
        await database.ensure_database_reachable(engine)
    assert "Name or service not known" in str(excinfo.value)


async def test_bad_credentials_are_not_masked_as_unavailable(no_async_sleep):
    # Retrying a wrong password would be pointless, and reporting it as
    # "database unavailable" would send the operator chasing the wrong problem.
    engine = _FakeEngine(
        _PGError("password authentication failed", "28P01"), then_succeed=False)
    with pytest.raises(_PGError):
        await database.ensure_database_reachable(engine)
    assert engine.calls == 1
    assert no_async_sleep == []