stackpilot/backend/tests/test_volume_prune.py

"""Volume prune with visible progress, and an explicit -a.

Same mechanism as the image prune (services/prune_job.py): the candidates are
listed, removed one by one and the progress is published through
/api/volumes/prune-status. On top of that the button used to call a bare
``docker volume prune``, which on Engine 23+ only removes anonymous volumes and
offered no way to include named ones.

Pinned here: nothing a container mounts is touched, named volumes only go with
-a, a refusal does not stop the rest, and the freed space is the sum of what
was actually deleted.
"""
from __future__ import annotations

import threading

import pytest
from docker.errors import APIError

MB = 1024 * 1024
ANON = {"com.docker.volume.anonymous": ""}
ANON_ID = "3f1c" + "0" * 60


class FakeApi:
    def __init__(self):
        self._volumes = [
            {"Name": ANON_ID, "Labels": ANON, "Scope": "local"},
            {"Name": "media_cache", "Labels": {"com.docker.compose.project": "media"}, "Scope": "local"},
            {"Name": "old_data", "Labels": None, "Scope": "local"},
            {"Name": "pg_data", "Labels": {}, "Scope": "local"},
            {"Name": "cluster_vol", "Labels": {}, "Scope": "global"},
        ]
        # pg_data is mounted by a (possibly stopped) container.
        self._containers = [
            {"Mounts": [{"Type": "volume", "Name": "pg_data"}, {"Type": "bind", "Source": "/srv"}]}
        ]
        self.removed: list[str] = []
        self.refuse: set[str] = set()
        self.gate: threading.Event | None = None

    def containers(self, all=False):
        return self._containers

    def volumes(self, filters=None):
        return {"Volumes": list(self._volumes), "Warnings": None}

    def remove_volume(self, name, force=False):
        assert force is False
        if self.gate:
            self.gate.wait(5)
        if name in self.refuse:
            raise APIError(f"remove {name}: volume is in use")
        self.removed.append(name)
        self._volumes = [v for v in self._volumes if v["Name"] != name]


class FakeClient:
    def __init__(self, api):
        self.api = api


@pytest.fixture
def svc(monkeypatch):
    from services import volume_service

    api = FakeApi()
    client = FakeClient(api)

    def short():
        raise AssertionError("prune must not use the 60s client")

    monkeypatch.setattr(volume_service, "get_client", short)
    monkeypatch.setattr(volume_service, "get_long_client", lambda: client)
    monkeypatch.setattr(
        volume_service,
        "volume_sizes",
        lambda force=False: {ANON_ID: 5 * MB, "media_cache": 300 * MB, "old_data": None, "pg_data": 900 * MB},
    )
    yield volume_service, api
    if volume_service._prune._lock.locked():
        volume_service._prune._lock.release()


def test_default_only_removes_unused_anonymous_volumes(svc):
    volume_service, api = svc

    result = volume_service.prune_volumes()

    assert api.removed == [ANON_ID]
    assert result == {"VolumesDeleted": [ANON_ID], "SpaceReclaimed": 5 * MB}


def test_all_also_removes_named_but_never_mounted_or_cluster_volumes(svc):
    volume_service, api = svc

    result = volume_service.prune_volumes(all_unused=True)

    assert set(api.removed) == {ANON_ID, "media_cache", "old_data"}
    # A volume of unknown size counts as nothing rather than failing the sum.
    assert result["SpaceReclaimed"] == 305 * MB
    status = volume_service.prune_status()
    assert (status["total"], status["done"], status["deleted"]) == (3, 3, 3)
    assert status["reclaimed"] == 305 * MB and status["active"] is False


def test_refused_volume_does_not_stop_the_rest(svc):
    volume_service, api = svc
    api.refuse.add("media_cache")

    result = volume_service.prune_volumes(all_unused=True)

    assert set(result["VolumesDeleted"]) == {ANON_ID, "old_data"}
    status = volume_service.prune_status()
    assert status["failed"] == 1 and status["errors"][0].startswith("media_cache: ")
    assert result["SpaceReclaimed"] == 5 * MB


def test_unknown_sizes_leave_reclaimed_unknown(svc, monkeypatch):
    volume_service, _ = svc

    def broken(force=False):
        raise RuntimeError("df failed")

    monkeypatch.setattr(volume_service, "volume_sizes", broken)
    result = volume_service.prune_volumes(all_unused=True)

    assert len(result["VolumesDeleted"]) == 3
    assert volume_service.prune_status()["reclaimed"] is None


def test_anonymous_volumes_are_labelled_by_short_id(svc):
    volume_service, api = svc
    api.gate = threading.Event()

    finished = threading.Event()

    volume_service.start_prune(all_unused=False, on_done=lambda r: finished.set())
    for _ in range(100):
        status = volume_service.prune_status()
        if status["current"]:
            break
        threading.Event().wait(0.02)
    api.gate.set()
    assert finished.wait(5)
    assert status["current"] == ANON_ID[:12]


def test_endpoint_background_then_poll(svc, as_admin):
    r = as_admin.post("/api/volumes/prune", params={"all": True, "background": True})
    assert r.status_code == 202 and r.json()["total"] == 3

    for _ in range(100):
        s = as_admin.get("/api/volumes/prune-status").json()
        if not s["active"]:
            break
        threading.Event().wait(0.02)
    assert s["deleted"] == 3 and s["reclaimed"] == 305 * MB


def test_endpoint_waits_by_default_and_is_audited(svc, as_admin):
    from sqlmodel import Session, select

    from database import engine
    from models.audit import AuditLog

    r = as_admin.post("/api/volumes/prune")
    assert r.status_code == 200
    assert r.json()["VolumesDeleted"] == [ANON_ID]

    with Session(engine) as session:
        row = session.exec(
            select(AuditLog).where(AuditLog.action == "volume.prune").order_by(AuditLog.id.desc())
        ).first()
    assert row.detail == f"all=False deleted=['{ANON_ID}']"


def test_endpoint_rejects_a_second_prune(svc, as_admin):
    volume_service, _ = svc
    volume_service._prune._lock.acquire()
    try:
        assert as_admin.post("/api/volumes/prune").status_code == 409
    finally:
        volume_service._prune._lock.release()


def test_prune_is_admin_only(svc, as_user):
    assert as_user.post("/api/volumes/prune").status_code == 403
    assert as_user.get("/api/volumes/prune-status").status_code == 403