tests/test_cli_sets.py

"""Command tests for set membership downloads and tag exports."""

from __future__ import annotations

import hashlib
import json
from collections.abc import Callable
from pathlib import Path

import httpx
import pytest
from typer.testing import CliRunner

import kscrape.cli.app as cli
from kscrape.sources.e6booru import E6BooruAdapter, RequestPacer
from kscrape.sources.models import SourceCredentials, SourceDefinition


@pytest.mark.parametrize("source", ["e6ai", "e621"])
def test_set_download_reads_all_members_without_post_search(
    source: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    """Use set membership even when a normal set search would omit a post."""

    requests: list[httpx.Request] = []
    file_bytes = b"example raster data"
    digest = hashlib.md5(file_bytes, usedforsecurity=False).hexdigest()

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        if request.url.path == "/post_sets.json":
            assert request.url.params["search[shortname]"] == "example_set"
            return httpx.Response(200, json=[_set_payload([8, 7])])
        if request.url.path in {"/posts/8.json", "/posts/7.json"}:
            post_id = int(request.url.path.split("/")[2].split(".")[0])
            return httpx.Response(
                200, json={"post": _post_payload(post_id, digest)}
            )
        if request.url.host == "files.example":
            return httpx.Response(200, content=file_bytes)
        raise AssertionError(f"Unexpected request: {request.url}")

    definition = SourceDefinition(source, source, f"https://{source}.example")
    adapter = E6BooruAdapter(
        definition,
        "kscrape-test/1.0",
        SourceCredentials(),
        client=httpx.Client(transport=httpx.MockTransport(handler)),
        pacer=RequestPacer(interval_seconds=0),
    )
    monkeypatch.setattr(cli, "_open_adapter", lambda *_args: adapter)

    result = CliRunner().invoke(
        cli.app,
        ["sets", "download", source, "example_set", "--output", str(tmp_path)],
    )

    assert result.exit_code == 0, result.output
    assert [request.url.path for request in requests] == [
        "/post_sets.json",
        "/posts/8.json",
        "/8.png",
        "/posts/7.json",
        "/7.png",
    ]
    set_root = tmp_path / "sets" / source / "example_set"
    assert (set_root / "images/8.png").read_bytes() == file_bytes
    assert (set_root / "images/7.png").read_bytes() == file_bytes
    records = [
        json.loads(line)
        for line in (set_root / "manifest.jsonl").read_text().splitlines()
    ]
    assert [record["relative_path"] for record in records] == [
        "images/8.png",
        "images/7.png",
    ]
    exported = json.loads(
        (tmp_path / "sets" / source / "example_set" / "tags.json").read_text()
    )
    assert exported["set"] == {
        "source": source,
        "id": 12,
        "name": "Example Set",
        "shortname": "example_set",
        "member_count": 2,
        "selected_count": 2,
    }
    assert [post["post_id"] for post in exported["posts"]] == [8, 7]
    assert exported["posts"][0]["tags"] == {"general": ["solo"]}
    assert exported["unavailable_post_ids"] == []


def test_set_download_accepts_id_limit_and_skips_existing_original(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    """Read a set by ID and select only the first member in set order."""

    image_bytes = b"existing image"
    digest = hashlib.md5(image_bytes, usedforsecurity=False).hexdigest()
    existing = tmp_path / "sets/e6ai/example_set/images/8.png"
    existing.parent.mkdir(parents=True)
    existing.write_bytes(image_bytes)
    requests: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request.url.path)
        if request.url.path == "/post_sets/12.json":
            return httpx.Response(200, json=_set_payload([8, 7]))
        if request.url.path == "/posts/8.json":
            return httpx.Response(200, json={"post": _post_payload(8, digest)})
        raise AssertionError(f"Unexpected request: {request.url}")

    adapter = _adapter(handler)
    monkeypatch.setattr(cli, "_open_adapter", lambda *_args: adapter)
    result = CliRunner().invoke(
        cli.app,
        [
            "sets",
            "download",
            "e6ai",
            "12",
            "--limit",
            "1",
            "--output",
            str(tmp_path),
        ],
    )

    assert result.exit_code == 0, result.output
    assert requests == ["/post_sets/12.json", "/posts/8.json"]
    manifest = json.loads(
        (tmp_path / "sets/e6ai/example_set/manifest.jsonl").read_text()
    )
    assert manifest["status"] == "skipped"
    exported = json.loads(
        (tmp_path / "sets/e6ai/example_set/tags.json").read_text()
    )
    assert exported["set"]["member_count"] == 2
    assert exported["set"]["selected_count"] == 1


def test_set_download_reports_unavailable_member_and_continues(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    """Record a missing member without losing later members or their tags."""

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path == "/post_sets.json":
            return httpx.Response(200, json=[_set_payload([8, 7])])
        if request.url.path == "/posts/8.json":
            return httpx.Response(404)
        if request.url.path == "/posts/7.json":
            return httpx.Response(200, json={"post": _post_payload(7, None)})
        if request.url.path == "/7.png":
            return httpx.Response(200, content=b"image")
        raise AssertionError(f"Unexpected request: {request.url}")

    adapter = _adapter(handler)
    monkeypatch.setattr(cli, "_open_adapter", lambda *_args: adapter)
    result = CliRunner().invoke(
        cli.app,
        ["sets", "download", "e6ai", "example_set", "--output", str(tmp_path)],
    )

    assert result.exit_code == 1
    assert "1 failed" in result.output
    records = [
        json.loads(line)
        for line in (tmp_path / "sets/e6ai/example_set/manifest.jsonl")
        .read_text()
        .splitlines()
    ]
    assert [record["status"] for record in records] == ["failed", "downloaded"]
    assert records[0]["relative_path"] is None
    exported = json.loads(
        (tmp_path / "sets/e6ai/example_set/tags.json").read_text()
    )
    assert exported["unavailable_post_ids"] == [8]
    assert [post["post_id"] for post in exported["posts"]] == [7]


def test_set_download_stops_on_access_error_without_complete_tag_export(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    """Preserve earlier files and manifest records if the source stops access."""

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path == "/post_sets.json":
            return httpx.Response(200, json=[_set_payload([8, 7])])
        if request.url.path == "/posts/8.json":
            return httpx.Response(200, json={"post": _post_payload(8, None)})
        if request.url.path == "/8.png":
            return httpx.Response(200, content=b"image")
        if request.url.path == "/posts/7.json":
            return httpx.Response(403)
        raise AssertionError(f"Unexpected request: {request.url}")

    monkeypatch.setattr(cli, "_open_adapter", lambda *_args: _adapter(handler))
    result = CliRunner().invoke(
        cli.app,
        ["sets", "download", "e6ai", "example_set", "--output", str(tmp_path)],
    )

    assert result.exit_code == 1
    assert "HTTP 403" in result.output
    assert (
        tmp_path / "sets/e6ai/example_set/images/8.png"
    ).read_bytes() == b"image"
    assert not (tmp_path / "sets/e6ai/example_set/manifest.jsonl").exists()
    assert not (tmp_path / "sets/e6ai/example_set/tags.json").exists()


def test_set_manifest_replaces_prior_attempts_on_rerun(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    """Keep one set record per member when a completed download is repeated."""

    image_bytes = b"image"
    digest = hashlib.md5(image_bytes, usedforsecurity=False).hexdigest()

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path == "/post_sets.json":
            return httpx.Response(200, json=[_set_payload([8, 7])])
        if request.url.path in {"/posts/8.json", "/posts/7.json"}:
            post_id = int(request.url.path.split("/")[2].split(".")[0])
            return httpx.Response(
                200, json={"post": _post_payload(post_id, digest)}
            )
        if request.url.host == "files.example":
            return httpx.Response(200, content=image_bytes)
        raise AssertionError(f"Unexpected request: {request.url}")

    arguments = [
        "sets",
        "download",
        "e6ai",
        "example_set",
        "--output",
        str(tmp_path),
    ]
    for _ in range(2):
        monkeypatch.setattr(
            cli, "_open_adapter", lambda *_args: _adapter(handler)
        )
        result = CliRunner().invoke(cli.app, arguments)
        assert result.exit_code == 0, result.output

    manifest_path = tmp_path / "sets/e6ai/example_set/manifest.jsonl"
    records = [
        json.loads(line) for line in manifest_path.read_text().splitlines()
    ]
    assert [record["post_id"] for record in records] == [8, 7]
    assert [record["status"] for record in records] == ["skipped", "skipped"]


def test_set_prune_removes_only_recorded_former_members(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    """Opt-in pruning leaves unrecorded files and other folders alone."""

    membership = [8, 7]
    image_bytes = b"image"
    digest = hashlib.md5(image_bytes, usedforsecurity=False).hexdigest()

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path == "/post_sets.json":
            return httpx.Response(200, json=[_set_payload(membership)])
        if request.url.path in {"/posts/8.json", "/posts/7.json"}:
            post_id = int(request.url.path.split("/")[2].split(".")[0])
            return httpx.Response(
                200, json={"post": _post_payload(post_id, digest)}
            )
        if request.url.host == "files.example":
            return httpx.Response(200, content=image_bytes)
        raise AssertionError(f"Unexpected request: {request.url}")

    arguments = [
        "sets",
        "download",
        "e6ai",
        "example_set",
        "--output",
        str(tmp_path),
    ]
    monkeypatch.setattr(cli, "_open_adapter", lambda *_args: _adapter(handler))
    first = CliRunner().invoke(cli.app, arguments)
    assert first.exit_code == 0, first.output
    set_root = tmp_path / "sets/e6ai/example_set"
    unrelated = set_root / "images/999.png"
    unrelated.write_bytes(b"unrelated")
    other_folder = tmp_path / "posts/e6ai/images/8.png"
    other_folder.parent.mkdir(parents=True)
    other_folder.write_bytes(b"post download")

    membership[:] = [7]
    second = CliRunner().invoke(cli.app, arguments)
    assert second.exit_code == 0, second.output
    assert (set_root / "images/8.png").exists()
    third = CliRunner().invoke(cli.app, [*arguments, "--prune"])

    assert third.exit_code == 0, third.output
    assert "Pruned 1" in third.output
    assert not (set_root / "images/8.png").exists()
    assert (set_root / "images/7.png").read_bytes() == image_bytes
    assert unrelated.read_bytes() == b"unrelated"
    assert other_folder.read_bytes() == b"post download"
    records = [
        json.loads(line)
        for line in (set_root / "manifest.jsonl").read_text().splitlines()
    ]
    assert [record["post_id"] for record in records] == [7]
    assert [
        entry["post_id"]
        for entry in json.loads((set_root / "image-index.json").read_text())
    ] == [7]


def test_set_prune_rejects_partial_selection(tmp_path: Path) -> None:
    """Never infer removed membership from a deliberately limited run."""

    result = CliRunner().invoke(
        cli.app,
        [
            "sets",
            "download",
            "e6ai",
            "example_set",
            "--output",
            str(tmp_path),
            "--limit",
            "1",
            "--prune",
        ],
    )

    assert result.exit_code == 2
    assert "complete set download" in result.output


def test_set_prune_waits_for_a_successful_run(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    """Retain former files and their ownership when a current member fails."""

    membership = [8]
    unavailable = False

    def handler(request: httpx.Request) -> httpx.Response:
        if request.url.path == "/post_sets.json":
            return httpx.Response(200, json=[_set_payload(membership)])
        if request.url.path == "/posts/7.json" and unavailable:
            return httpx.Response(404)
        if request.url.path in {"/posts/8.json", "/posts/7.json"}:
            post_id = int(request.url.path.split("/")[2].split(".")[0])
            return httpx.Response(
                200, json={"post": _post_payload(post_id, None)}
            )
        if request.url.host == "files.example":
            return httpx.Response(200, content=b"image")
        raise AssertionError(f"Unexpected request: {request.url}")

    monkeypatch.setattr(cli, "_open_adapter", lambda *_args: _adapter(handler))
    arguments = [
        "sets",
        "download",
        "e6ai",
        "example_set",
        "--output",
        str(tmp_path),
    ]
    assert CliRunner().invoke(cli.app, arguments).exit_code == 0
    old_image = tmp_path / "sets/e6ai/example_set/images/8.png"

    membership[:] = [7]
    unavailable = True
    failed = CliRunner().invoke(cli.app, [*arguments, "--prune"])
    assert failed.exit_code == 1
    assert old_image.read_bytes() == b"image"

    unavailable = False
    completed = CliRunner().invoke(cli.app, [*arguments, "--prune"])
    assert completed.exit_code == 0, completed.output
    assert not old_image.exists()


def _adapter(
    handler: Callable[[httpx.Request], httpx.Response],
) -> E6BooruAdapter:
    """Wrap a mock transport with the real adapter and no test delay.

    :param handler: Function that returns mocked source responses.
    :return: Adapter that exercises real set, post, and file request paths.
    """

    return E6BooruAdapter(
        SourceDefinition("e6ai", "e6AI", "https://e6ai.example"),
        "kscrape-test/1.0",
        SourceCredentials(),
        client=httpx.Client(transport=httpx.MockTransport(handler)),
        pacer=RequestPacer(interval_seconds=0),
    )


def _set_payload(post_ids: list[int]) -> dict[str, object]:
    """Return the source's set record with ordered member IDs.

    :param post_ids: IDs to expose as this set's membership.
    :return: JSON-ready set metadata.
    """

    return {
        "id": 12,
        "name": "Example Set",
        "shortname": "example_set",
        "post_count": len(post_ids),
        "post_ids": post_ids,
    }


def _post_payload(post_id: int, md5: str | None) -> dict[str, object]:
    """Return file and tag fields required by the source parser.

    :param post_id: Member identifier used by the direct post endpoint.
    :param md5: Optional original checksum for skip verification.
    :return: JSON-ready post metadata.
    """

    return {
        "id": post_id,
        "rating": "s",
        "file": {
            "ext": "png",
            "url": f"https://files.example/{post_id}.png",
            "md5": md5,
        },
        "tags": {"general": ["solo"]},
    }