stackpilot/backend/tests/test_compose_runs.py

"""Streamed compose runs: the deploy console and the update bar.

Until 0.67.1 the run *was* the WebSocket handler, so closing the console in the
middle of a deploy released the stack lock at once — while compose went on —
and logged the deploy as a success with ``rc=None``. The run is its own task
now; a socket only watches it. And the deploy socket is admin-only, like the
REST ``/start`` it mirrors: before, the read-only role could run ``compose up``
through it.
"""
from __future__ import annotations

import asyncio
import json

import pytest
from sqlmodel import Session


@pytest.fixture(autouse=True)
def clean_locks(db):
    from sqlmodel import delete

    from database import engine
    from models.runtime_state import StackLock

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


def _locked(stack_id: str) -> bool:
    from database import engine
    from services import stack_lock_service

    with Session(engine) as session:
        return stack_lock_service.is_busy(session, stack_id)


def test_the_lock_is_held_until_compose_finishes(db):
    from services import compose_run_service as runs

    async def scenario():
        release = asyncio.Event()
        finished: list = []

        async def stream():
            yield ("log", "pulling")
            await release.wait()
            yield ("log", "started")
            yield ("done", 0)

        async def on_finish(rc):
            finished.append(rc)

        run, queue, attached = runs.start("rx", "start", "alice", stream, on_finish)
        assert not attached
        assert await queue.get() == ("log", "pulling")
        # Nobody is watching any more — the console was closed.
        run.unsubscribe(queue)
        await asyncio.sleep(0.05)
        assert _locked("rx"), "the lock must outlive the socket"
        assert finished == []

        release.set()
        for _ in range(50):
            if finished:
                break
            await asyncio.sleep(0.02)
        assert finished == [0], "the real exit code, after the fact"
        assert not _locked("rx")

    asyncio.run(scenario())


def test_a_second_socket_attaches_to_the_run_in_flight(db):
    from services import compose_run_service as runs

    async def scenario():
        release = asyncio.Event()

        async def stream():
            yield ("log", "one")
            await release.wait()
            yield ("done", 0)

        run, first, _ = runs.start("ry", "start", "alice", stream)
        assert await first.get() == ("log", "one")

        again, second, attached = runs.start("ry", "start", "bob", stream)
        assert attached and again is run
        assert await second.get() == ("log", "one"), "backlog first"

        with pytest.raises(runs.RunBusy):
            runs.start("ry", "update", "carol", stream)

        release.set()
        items = []
        while (item := await second.get()) is not None:
            items.append(item)
        assert items == [("done", 0)]

    asyncio.run(scenario())


def test_a_run_that_cannot_get_the_lock_is_busy(db):
    from database import engine
    from services import compose_run_service as runs
    from services import stack_lock_service

    with Session(engine) as session:
        stack_lock_service.acquire(session, "rz", "backup", "scheduled")

    async def stream():
        yield ("done", 0)

    async def scenario():
        with pytest.raises(runs.RunBusy) as excinfo:
            runs.start("rz", "start", "alice", stream)
        assert excinfo.value.action == "backup"

    asyncio.run(scenario())


def test_a_failing_stream_still_releases_the_lock(db):
    from services import compose_run_service as runs

    async def scenario():
        async def stream():
            yield ("log", "about to break")
            raise RuntimeError("compose binary missing")

        finished = []

        async def on_finish(rc):
            finished.append(rc)

        _run, queue, _ = runs.start("rf", "update", "alice", stream, on_finish)
        items = []
        while (item := await queue.get()) is not None:
            items.append(item)
        await asyncio.sleep(0.05)
        assert items[-1][0] == "done" and items[-1][1] != 0
        assert finished and finished[0] != 0
        assert not _locked("rf")

    asyncio.run(scenario())


def test_the_deploy_socket_is_admin_only(client, user_token):
    from starlette.websockets import WebSocketDisconnect

    with pytest.raises(WebSocketDisconnect) as excinfo:
        with client.websocket_connect(f"/ws/deploy/anything?token={user_token}") as socket:
            socket.receive_text()
    assert excinfo.value.code == 4403


def test_deploying_an_unknown_stack_is_refused(client, admin_token):
    from starlette.websockets import WebSocketDisconnect

    with pytest.raises(WebSocketDisconnect) as excinfo:
        with client.websocket_connect(f"/ws/deploy/no-such-stack?token={admin_token}") as socket:
            frame = json.loads(socket.receive_text())
            assert frame["type"] == "error"
            socket.receive_text()
    assert excinfo.value.code == 4404


# --------------------------------------------------------------------------- #
# The REST paths
# --------------------------------------------------------------------------- #


def test_update_stops_at_a_failed_pull(monkeypatch):
    """`up` after a failed pull recreated nothing and reported success."""
    from services import compose_service

    calls = []

    async def fake(stack_id, args, override=None, timeout=600.0):
        calls.append(args[0])
        rc = 1 if args[0] == "pull" else 0
        return {"returncode": rc, "stdout": "", "stderr": "denied", "command": args[0]}

    monkeypatch.setattr(compose_service, "run_compose", fake)
    result = asyncio.run(compose_service.update("s"))
    assert calls == ["pull"]
    assert result["returncode"] == 1


def test_auto_update_reports_a_failed_compose_run(db, monkeypatch):
    """The exit codes used to be ignored: a failure was recorded as "updated"."""
    from database import engine
    from models.auto_update import AutoUpdate
    from services import auto_update_service, compose_service, update_service

    monkeypatch.setattr(compose_service, "compute_status", lambda stack_id: "running")

    async def stale(stack_id, refresh=False):
        return {"update_available": True, "stale_images": ["nginx:latest"]}

    async def failing_update(stack_id, override=None):
        return {"returncode": 1, "stdout": "", "stderr": "Error: pull access denied"}

    notified = []

    async def notify(event, title, message, session=None):
        notified.append(title)

    monkeypatch.setattr(update_service, "stack_updates", stale)
    monkeypatch.setattr(compose_service, "update", failing_update)
    monkeypatch.setattr(auto_update_service.notify_service, "notify", notify)

    with Session(engine) as session:
        policy = AutoUpdate(stack_id="au-fail", enabled=True, redeploy=True)
        session.add(policy)
        session.commit()
        session.refresh(policy)
        asyncio.run(auto_update_service.run_policy(session, policy))
        session.refresh(policy)
        assert policy.last_status == "error"
        assert "pull access denied" in (policy.last_result or "")
        session.delete(policy)
        session.commit()
    assert notified and "failed" in notified[0]