stackpilot/backend/tests/test_runtime_state.py

"""State that used to live in module dicts (F10, F11).

Three things were kept in process memory, all of them assuming exactly one
uvicorn worker without saying so, and all of them lost on restart:

* the live stats sweep, which was not cached at all and got polled every five
  seconds by two pages at once;
* the registry digests behind the update badges, so a restart blanked every
  badge for up to an hour and re-announced updates already notified about;
* the login rate limiter, which an attacker could reset by getting the process
  to restart, and which multiplied by the worker count.
"""
from __future__ import annotations

from datetime import datetime, timedelta, timezone

import pytest
from sqlmodel import Session, delete, select


# --------------------------------------------------------------------------- #
# F10 — the stats cache
# --------------------------------------------------------------------------- #


@pytest.fixture
def stats(monkeypatch):
    from services import stats_service

    stats_service._cache["data"] = None
    stats_service._cache["ts"] = 0.0
    return stats_service


def test_repeat_calls_inside_the_window_reuse_one_sweep(stats, monkeypatch):
    """The dashboard and the stacks list both poll this every 5s.

    Each sweep is one blocking daemon call per running container, so serving
    them from one sample is the whole point.
    """
    calls = []
    monkeypatch.setattr(stats, "_sample", lambda: calls.append(1) or {"a": {}})

    assert stats.stack_stats() == {"a": {}}
    assert stats.stack_stats() == {"a": {}}
    assert stats.stack_stats() == {"a": {}}
    assert len(calls) == 1


def test_the_cache_expires(stats, monkeypatch):
    calls = []
    monkeypatch.setattr(stats, "_sample", lambda: calls.append(1) or {})

    stats.stack_stats()
    # Pretend the window has passed rather than sleeping through it.
    stats._cache["ts"] -= stats.CACHE_TTL + 1
    stats.stack_stats()
    assert len(calls) == 2


def test_refresh_bypasses_the_cache(stats, monkeypatch):
    calls = []
    monkeypatch.setattr(stats, "_sample", lambda: calls.append(1) or {})

    stats.stack_stats()
    stats.stack_stats(refresh=True)
    assert len(calls) == 2


def test_a_failing_sample_is_not_cached_as_a_success(stats, monkeypatch):
    """A Docker outage returns {} — it must not stick for the whole window
    once the daemon is back."""
    monkeypatch.setattr(stats, "_sample", dict)
    assert stats.stack_stats() == {}

    monkeypatch.setattr(stats, "_sample", lambda: {"a": {"containers": 1}})
    assert stats.stack_stats(refresh=True) == {"a": {"containers": 1}}


# --------------------------------------------------------------------------- #
# F11 — the image update cache survives a restart
# --------------------------------------------------------------------------- #


@pytest.fixture
def clean_image_status(db):
    from database import engine
    from models.runtime_state import ImageStatus
    from services import update_service

    with Session(engine) as session:
        session.exec(delete(ImageStatus))
        session.commit()
    update_service._CACHE.clear()
    update_service._NOTIFIED.clear()
    update_service.set_persist_callback(None)
    yield
    update_service.set_persist_callback(None)
    update_service._CACHE.clear()
    update_service._NOTIFIED.clear()


def test_a_saved_status_comes_back_after_a_restart(clean_image_status):
    """The badge is there immediately instead of blank until the next sweep."""
    from services import image_status_store, update_service

    status = update_service.UpdateStatus(
        image="nginx:latest",
        update_available=True,
        current_digest="sha256:old",
        remote_digest="sha256:new",
        checked_at=123.0,
    )
    image_status_store.save(status, notified=True)

    # Simulate a restart: memory is empty, then install() seeds it.
    update_service._CACHE.clear()
    update_service._NOTIFIED.clear()
    assert image_status_store.install() == 1

    cached = update_service.get_cache()["nginx:latest"]
    assert cached["update_available"] is True
    assert cached["remote_digest"] == "sha256:new"
    # And the "already told the user" mark comes back with it, so the restart
    # does not re-announce the same update.
    assert "nginx:latest" in update_service._NOTIFIED


def test_installing_starts_mirroring_writes(clean_image_status):
    from database import engine
    from models.runtime_state import ImageStatus
    from services import image_status_store, update_service

    image_status_store.install()
    update_service._store(
        update_service.UpdateStatus(
            image="redis:7", update_available=False, current_digest="sha256:a",
            remote_digest="sha256:a", checked_at=1.0,
        ),
        notified=False,
    )
    with Session(engine) as session:
        row = session.get(ImageStatus, "redis:7")
    assert row is not None and row.update_available is False


def test_saving_the_same_image_twice_updates_it(clean_image_status):
    from database import engine
    from models.runtime_state import ImageStatus
    from services import image_status_store, update_service

    def status(remote):
        return update_service.UpdateStatus(
            image="app:1", update_available=True, current_digest="sha256:x",
            remote_digest=remote, checked_at=1.0,
        )

    image_status_store.save(status("sha256:one"), notified=False)
    image_status_store.save(status("sha256:two"), notified=True)

    with Session(engine) as session:
        rows = session.exec(select(ImageStatus).where(ImageStatus.image == "app:1")).all()
    assert len(rows) == 1
    assert rows[0].remote_digest == "sha256:two"
    assert rows[0].notified is True


def test_pruning_drops_images_no_stack_uses_any_more(clean_image_status):
    from database import engine
    from models.runtime_state import ImageStatus
    from services import image_status_store, update_service

    for image in ("kept:1", "gone:1"):
        image_status_store.save(
            update_service.UpdateStatus(
                image=image, update_available=False, current_digest=None,
                remote_digest=None, checked_at=0.0,
            ),
            notified=False,
        )

    assert image_status_store.prune(keep={"kept:1"}) == 1
    with Session(engine) as session:
        assert [r.image for r in session.exec(select(ImageStatus)).all()] == ["kept:1"]


def test_the_agent_has_no_persistence_wired_up(clean_image_status):
    """``update_service`` is shared with the agent, which has no database.

    Persistence therefore has to be opt-in, registered by the central app —
    if this module ever imports the database directly, the agent breaks.
    """
    import inspect

    from services import update_service

    source = inspect.getsource(update_service)
    assert "from database import" not in source
    assert "import database" not in source
    assert update_service._persist_cb is None


# --------------------------------------------------------------------------- #
# F11 — the login rate limiter
# --------------------------------------------------------------------------- #


@pytest.fixture
def clean_attempts(db):
    from database import engine
    from models.runtime_state import LoginAttempt

    with Session(engine) as session:
        session.exec(delete(LoginAttempt))
        session.commit()
    yield


def test_the_eleventh_attempt_in_a_minute_is_refused(client, clean_attempts):
    body = {"username": "nobody", "password": "wrong"}
    for _ in range(10):
        assert client.post("/api/auth/login", json=body).status_code == 401
    assert client.post("/api/auth/login", json=body).status_code == 429


def test_the_limit_survives_a_restart(db, clean_attempts):
    """In memory this reset on every restart, so an attacker could clear their
    own budget by getting the process to fall over."""
    from database import engine
    from fastapi import HTTPException
    from routers import auth as auth_router

    with Session(engine) as session:
        for _ in range(10):
            auth_router._check_rate_limit(session, "10.0.0.9")

    # A fresh session stands in for a fresh process — the counter is a table.
    with Session(engine) as session:
        with pytest.raises(HTTPException) as excinfo:
            auth_router._check_rate_limit(session, "10.0.0.9")
    assert excinfo.value.status_code == 429


def test_different_clients_have_separate_budgets(db, clean_attempts):
    """Only meaningful because uvicorn runs with --proxy-headers; without it
    every request carries the frontend container's IP and one client would
    throttle everyone."""
    from database import engine
    from fastapi import HTTPException
    from routers import auth as auth_router

    with Session(engine) as session:
        for _ in range(10):
            auth_router._check_rate_limit(session, "10.0.0.1")
        auth_router._check_rate_limit(session, "10.0.0.2")  # must not raise
        with pytest.raises(HTTPException):
            auth_router._check_rate_limit(session, "10.0.0.1")


def test_attempts_age_out_of_the_window(db, clean_attempts):
    from database import engine
    from models.runtime_state import LoginAttempt
    from routers import auth as auth_router

    old = datetime.now(timezone.utc) - timedelta(minutes=5)
    with Session(engine) as session:
        for _ in range(10):
            session.add(LoginAttempt(ip="10.0.0.3", at=old))
        session.commit()
        auth_router._check_rate_limit(session, "10.0.0.3")  # window has moved on


def test_stale_rows_are_pruned(db, clean_attempts):
    from database import engine
    from models.runtime_state import LoginAttempt
    from routers import auth as auth_router

    ancient = datetime.now(timezone.utc) - timedelta(days=2)
    with Session(engine) as session:
        session.add(LoginAttempt(ip="10.0.0.4", at=ancient))
        session.commit()
        auth_router._check_rate_limit(session, "10.0.0.5")
        remaining = session.exec(select(LoginAttempt)).all()
    assert all(r.ip != "10.0.0.4" for r in remaining), "the table must not grow forever"