stackpilot/backend/tests/test_image_prune.py

"""Image prune with visible progress.

``docker image prune -a`` is one blocking call that reports nothing until it
is done; on a host with a lot of images that meant minutes of a spinner and
then, past the proxy timeout, an error while the daemon carried on. The prune
now removes the candidates one by one, publishes how far it is through
/api/images/prune-status, and the UI starts it in the background and polls.

What these tests pin: the candidate set matches Docker's (in-use images are
never touched, dangling-only without -a), multi-tag images lose every tag,
one failing image does not stop the rest, reclaimed space is measured rather
than summed, and the endpoints keep the old synchronous contract for scripts.
"""
from __future__ import annotations

import threading

import pytest
from docker.errors import APIError

MB = 1024 * 1024


class FakeApi:
    """Just enough of docker-py's low-level APIClient for the prune."""

    def __init__(self, images, containers, layers=500 * MB):
        self._images = images
        self._containers = containers
        self.layers = layers
        self.removed: list[str] = []
        self.refuse: set[str] = set()
        self.gate: threading.Event | None = None

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

    def images(self, filters=None):
        if filters and filters.get("dangling"):
            return [i for i in self._images if not i.get("RepoTags")]
        return list(self._images)

    def remove_image(self, ref, force=False, noprune=False):
        if self.gate:
            self.gate.wait(5)
        if ref in self.refuse:
            raise APIError(f"conflict: unable to remove {ref}")
        self.removed.append(ref)
        for img in self._images:
            if ref in (img.get("RepoTags") or []):
                img["RepoTags"].remove(ref)
                if img["RepoTags"]:
                    return [{"Untagged": ref}]
                self._images.remove(img)
                self.layers -= img["Size"]
                return [{"Untagged": ref}, {"Deleted": img["Id"]}]
            if img["Id"] == ref:
                self._images.remove(img)
                self.layers -= img["Size"]
                return [{"Deleted": img["Id"]}]
        raise APIError(f"No such image: {ref}")

    # /system/df?type=image
    def _url(self, path):
        return path

    def _get(self, url, params=None):
        assert params == {"type": "image"}
        return self

    def _result(self, response, json=False):
        return {"LayersSize": self.layers}


class FakeImages:
    def __init__(self, api):
        self.api = api
        self.prune_filters = None

    def prune(self, filters=None):
        self.prune_filters = filters
        return {"ImagesDeleted": None, "SpaceReclaimed": 0}


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


def _host():
    images = [
        {"Id": "sha256:aaa", "RepoTags": ["nginx:1.25"], "Size": 100 * MB},
        {"Id": "sha256:bbb", "RepoTags": ["redis:7", "redis:latest"], "Size": 50 * MB},
        {"Id": "sha256:ccc", "RepoTags": [], "Size": 20 * MB},
        {"Id": "sha256:ddd", "RepoTags": ["postgres:16"], "Size": 200 * MB},
    ]
    # postgres is used by a (possibly stopped) container and must survive.
    containers = [{"ImageID": "sha256:ddd"}]
    return FakeApi(images, containers)


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

    api = _host()
    client = FakeClient(api)

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

    monkeypatch.setattr(image_service, "get_client", short)
    monkeypatch.setattr(image_service, "get_long_client", lambda: client)
    yield image_service, api, client
    if image_service._prune._lock.locked():  # a test that failed mid-prune
        image_service._prune._lock.release()


def test_all_unused_removes_everything_no_container_uses(svc):
    image_service, api, client = svc

    result = image_service.prune_images(all_unused=True)

    assert {i["Id"] for i in api._images} == {"sha256:ddd"}
    # Every tag of a multi-tag image is removed, without force.
    assert set(api.removed) == {"nginx:1.25", "redis:7", "redis:latest", "sha256:ccc"}
    assert {e["Deleted"] for e in result["ImagesDeleted"] if "Deleted" in e} == {
        "sha256:aaa", "sha256:bbb", "sha256:ccc",
    }
    assert result["SpaceReclaimed"] == 170 * MB
    # The daemon's own prune still runs afterwards to sweep up leftovers.
    assert client.images.prune_filters == {"dangling": False}

    status = image_service.prune_status()
    assert status["active"] is False
    assert (status["total"], status["done"], status["deleted"]) == (3, 3, 3)
    assert status["reclaimed"] == 170 * MB


def test_default_only_touches_dangling_images(svc):
    image_service, api, client = svc

    image_service.prune_images(all_unused=False)

    assert api.removed == ["sha256:ccc"]
    assert client.images.prune_filters == {"dangling": True}


def test_one_refused_image_does_not_stop_the_rest(svc):
    image_service, api, _ = svc
    api.refuse.add("nginx:1.25")

    image_service.prune_images(all_unused=True)

    status = image_service.prune_status()
    assert status["failed"] == 1 and status["done"] == 3 and status["deleted"] == 2
    assert status["errors"][0].startswith("nginx:1.25: ")
    assert any(i["Id"] == "sha256:aaa" for i in api._images)


def test_background_prune_reports_progress_and_blocks_a_second_one(svc):
    image_service, api, _ = svc
    api.gate = threading.Event()
    finished = threading.Event()

    started = image_service.start_prune(True, on_done=lambda r: finished.set())
    assert started["active"] is True and started["total"] == 3
    from services.prune_job import PruneBusy

    with pytest.raises(PruneBusy):
        image_service.start_prune(True)

    api.gate.set()
    assert finished.wait(5)
    status = image_service.prune_status()
    assert status["run"] == started["run"]
    assert status["done"] == 3 and status["reclaimed"] == 170 * MB


def test_background_prune_error_lands_in_the_status(svc):
    from docker_client import DockerError

    image_service, _, client = svc

    def broken(filters=None):
        raise DockerError("docker_timeout", "Docker did not answer in time")

    client.images.prune = broken
    image_service.start_prune(True)
    for _ in range(100):
        if not image_service.prune_status()["active"]:
            break
        threading.Event().wait(0.02)
    status = image_service.prune_status()
    assert status["active"] is False
    assert status["error"] == "Docker did not answer in time"
    assert not image_service._prune._lock.locked()


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/images/prune", params={"all": True})
    assert r.status_code == 200
    assert r.json()["SpaceReclaimed"] == 170 * MB

    with Session(engine) as session:
        row = session.exec(
            select(AuditLog).where(AuditLog.action == "image.prune").order_by(AuditLog.id.desc())
        ).first()
    assert row.detail == f"all=True deleted=3 reclaimed={170 * MB}"


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

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


def test_endpoint_rejects_a_second_prune(svc, as_admin):
    image_service, _, _ = svc
    image_service._prune._lock.acquire()
    try:
        r = as_admin.post("/api/images/prune", params={"background": True})
        assert r.status_code == 409
    finally:
        image_service._prune._lock.release()


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