AIChatExporter/tests/test_media.py

"""Tests for media downloads and Joplin resource rewriting."""

from pathlib import Path

import pytest

from src.blocks import (
    make_file_placeholder,
    make_image_placeholder,
    render_blocks_to_markdown,
)
from src.loss_report import LossReport
from src.media import resolve_media, resolve_media_policy
from src.providers.base import ProviderError
from src.providers.chatgpt import parse_asset_file_id


# ---------------------------------------------------------------------------
# Asset reference parsing
# ---------------------------------------------------------------------------


class TestParseAssetFileId:
    def test_plain_sediment(self):
        assert parse_asset_file_id("sediment://file_00000000245c71fda5") == "file_00000000245c71fda5"

    def test_generated_image_with_hash_and_page(self):
        ref = "sediment://8456107fc383a53#file_00000000979c71f685#p_6.png"
        assert parse_asset_file_id(ref) == "file_00000000979c71f685"

    def test_file_service_scheme(self):
        assert parse_asset_file_id("file-service://file-AbCdEf") == "file-AbCdEf"

    def test_unrecognised(self):
        assert parse_asset_file_id("https://example.com/x.png") is None
        assert parse_asset_file_id("") is None
        assert parse_asset_file_id(None) is None


# ---------------------------------------------------------------------------
# Media policy
# ---------------------------------------------------------------------------


class TestMediaPolicy:
    def test_default(self, monkeypatch):
        monkeypatch.delenv("EXPORTER_DOWNLOAD_MEDIA", raising=False)
        assert resolve_media_policy() == "images"

    def test_valid(self, monkeypatch):
        monkeypatch.setenv("EXPORTER_DOWNLOAD_MEDIA", "all")
        assert resolve_media_policy() == "all"

    def test_invalid_falls_back(self, monkeypatch, caplog):
        monkeypatch.setenv("EXPORTER_DOWNLOAD_MEDIA", "bogus")
        assert resolve_media_policy() == "images"


# ---------------------------------------------------------------------------
# resolve_media
# ---------------------------------------------------------------------------


class _FakeProvider:
    """Minimal provider exposing download_asset + the ref parser."""

    def __init__(self, assets=None, fail_refs=None):
        self._assets = assets or {}
        self._fail_refs = fail_refs or {}
        self.calls = []

    def parse_asset_file_id(self, ref):
        return parse_asset_file_id(ref)

    def download_asset(self, ref):
        self.calls.append(ref)
        if ref in self._fail_refs:
            raise ProviderError("chatgpt", "download_asset", self._fail_refs[ref])
        return self._assets[ref]  # (content, mime, file_name)


def _conv_with(blocks):
    return {
        "id": "conv-1",
        "title": "Has Media",
        "provider": "chatgpt",
        "project": None,
        "created_at": "2026-05-20T00:00:00+00:00",
        "messages": [{"role": "user", "blocks": blocks}],
    }


class TestResolveMedia:
    def test_downloads_image_and_inlines(self, tmp_path):
        ref = "sediment://file_img1"
        provider = _FakeProvider({ref: (b"\x89PNG\r\n", "image/png", "x.png")})
        block = make_image_placeholder(ref=ref, source="user_upload")
        conv = _conv_with([block])
        report = LossReport()

        n = resolve_media(conv, provider, tmp_path, "provider/project/year", "images", report)

        assert n.downloaded == 1
        assert report.media_downloaded == 1
        assert block["local_path"] == "media/file_img1.png"
        rendered = render_blocks_to_markdown([block])
        assert rendered == "![user_upload](media/file_img1.png)"
        # File written under the conversation's media/ dir
        written = list(tmp_path.rglob("media/file_img1.png"))
        assert written and written[0].read_bytes() == b"\x89PNG\r\n"

    def test_images_policy_skips_files(self, tmp_path):
        ref = "sediment://file_audio1"
        provider = _FakeProvider({ref: (b"RIFF", "audio/wav", "a.wav")})
        block = make_file_placeholder(ref=ref, mime="audio/wav", size_bytes=1000)
        conv = _conv_with([block])
        report = LossReport()

        n = resolve_media(conv, provider, tmp_path, "provider/project/year", "images", report)
        assert n.downloaded == 0
        assert "local_path" not in block
        assert provider.calls == []

    def test_all_policy_downloads_files(self, tmp_path):
        ref = "sediment://file_audio1"
        provider = _FakeProvider({ref: (b"RIFFdata", "audio/wav", "a.wav")})
        block = make_file_placeholder(ref=ref, mime="audio/wav", size_bytes=8)
        conv = _conv_with([block])
        report = LossReport()

        n = resolve_media(conv, provider, tmp_path, "provider/project/year", "all", report)
        assert n.downloaded == 1
        assert block["local_path"] == "media/file_audio1.wav"
        rendered = render_blocks_to_markdown([block])
        assert "media/file_audio1.wav" in rendered and rendered.startswith("> 📎")

    def test_off_policy_noop(self, tmp_path):
        provider = _FakeProvider({"sediment://file_x": (b"x", "image/png", None)})
        block = make_image_placeholder(ref="sediment://file_x", source="user_upload")
        conv = _conv_with([block])
        report = LossReport()
        assert resolve_media(
            conv, provider, tmp_path, "provider/project/year", "off", report
        ).downloaded == 0
        assert provider.calls == []

    def test_idempotent_uses_disk(self, tmp_path):
        ref = "sediment://file_img1"
        provider = _FakeProvider({ref: (b"\x89PNG", "image/png", "x.png")})
        conv = _conv_with([make_image_placeholder(ref=ref, source="user_upload")])
        report = LossReport()
        resolve_media(conv, provider, tmp_path, "provider/project/year", "images", report)
        assert len(provider.calls) == 1

        # Second run, fresh blocks: file already on disk → no new API call.
        conv2 = _conv_with([make_image_placeholder(ref=ref, source="user_upload")])
        resolve_media(conv2, provider, tmp_path, "provider/project/year", "images", LossReport())
        assert len(provider.calls) == 1  # unchanged
        assert conv2["messages"][0]["blocks"][0]["local_path"] == "media/file_img1.png"

    def test_failure_keeps_placeholder_and_counts(self, tmp_path):
        ref = "sediment://file_gone"
        provider = _FakeProvider(fail_refs={ref: RuntimeError("Signed URL returned HTTP 404")})
        block = make_image_placeholder(ref=ref, source="model_generated")
        conv = _conv_with([block])
        report = LossReport()

        n = resolve_media(conv, provider, tmp_path, "provider/project/year", "images", report)
        assert n.downloaded == 0
        assert "local_path" not in block
        assert report.media_failed["expired-or-missing"] == 1
        # Still renders as a placeholder, not a broken image link
        assert render_blocks_to_markdown([block]).startswith("> 🖼️")

    def test_forbidden_counted_separately_from_generic_error(self, tmp_path):
        """403 is a distinct bucket: the asset exists, we were refused."""
        ref = "sediment://file_denied"
        provider = _FakeProvider(
            fail_refs={ref: RuntimeError("HTTP 403 — detail: unauthorized")}
        )
        block = make_image_placeholder(ref=ref, source="model_generated")
        report = LossReport()

        resolve_media(
            _conv_with([block]), provider, tmp_path, "provider/project/year",
            "images", report,
        )
        assert report.media_failed["forbidden"] == 1
        assert "download-error" not in report.media_failed

    def test_provider_without_download_asset(self, tmp_path):
        """claude-code has no remote assets — resolve_media must no-op."""
        class NoDownload:
            pass
        block = make_image_placeholder(ref="sediment://file_x", source="user_upload")
        conv = _conv_with([block])
        assert resolve_media(
            conv, NoDownload(), tmp_path, "provider/project/year", "images", LossReport()
        ).downloaded == 0


# ---------------------------------------------------------------------------
# Joplin media rewriting
# ---------------------------------------------------------------------------


class _FakeJoplin:
    def __init__(self):
        self.uploaded = []
        self._n = 0

    def create_resource(self, file_path, title=None):
        self.uploaded.append(Path(file_path).name)
        self._n += 1
        return f"res{self._n}"


class TestUploadMediaAndRewrite:
    def _note(self, tmp_path, body):
        media = tmp_path / "media"
        media.mkdir()
        (media / "file_img1.png").write_bytes(b"\x89PNG")
        (media / "clip.wav").write_bytes(b"RIFF")
        return body

    def test_rewrites_image_and_file_links(self, tmp_path):
        from src.joplin import upload_media_and_rewrite
        body = self._note(
            tmp_path,
            "Look: ![user_upload](media/file_img1.png)\n"
            "> 📎 **File attached** — [clip.wav](media/clip.wav) (audio/wav)",
        )
        client = _FakeJoplin()
        new_body, res_map = upload_media_and_rewrite(body, tmp_path, client, {})

        assert "![user_upload](:/res1)" in new_body
        assert "(:/res2)" in new_body
        assert "media/" not in new_body
        assert res_map == {"media/file_img1.png": "res1", "media/clip.wav": "res2"}
        assert len(client.uploaded) == 2

    def test_reuses_known_resource_ids(self, tmp_path):
        from src.joplin import upload_media_and_rewrite
        body = self._note(tmp_path, "![x](media/file_img1.png)")
        client = _FakeJoplin()
        existing = {"media/file_img1.png": "existing-res"}
        new_body, res_map = upload_media_and_rewrite(body, tmp_path, client, existing)
        assert "(:/existing-res)" in new_body
        assert client.uploaded == []  # no re-upload
        assert res_map == existing

    def test_missing_file_left_as_link(self, tmp_path):
        from src.joplin import upload_media_and_rewrite
        (tmp_path / "media").mkdir()
        body = "![x](media/gone.png)"
        client = _FakeJoplin()
        new_body, res_map = upload_media_and_rewrite(body, tmp_path, client, {})
        assert new_body == body
        assert res_map == {}
        assert client.uploaded == []

    def test_no_media_links_untouched(self, tmp_path):
        from src.joplin import upload_media_and_rewrite
        body = "Just text with [a link](https://example.com) and `media/foo` in code."
        client = _FakeJoplin()
        new_body, res_map = upload_media_and_rewrite(body, tmp_path, client, {})
        assert new_body == body
        assert res_map == {}


# ---------------------------------------------------------------------------
# Transient vs permanent failures, and the cross-run retry they feed.
#
# Measured 2026-10-20 in one scheduled run: a signed URL returned 500 on one
# asset and 503 on another. Both are recoverable, yet the conversation was
# marked exported, so the assets were lost for good. Permanent failures must
# stay out of the retry set — 19 assets in this archive 403/404 because their
# bytes were lost in a 2026-07-14 Library migration, and retrying those every
# morning would be noise that also trips the circuit breaker.
# ---------------------------------------------------------------------------


class TestTransientFailureRetry:
    def _run(self, failure, tmp_path, circuit=None):
        from src.media import resolve_media

        ref = "sediment://file_x"
        provider = _FakeProvider(fail_refs={ref: RuntimeError(failure)})
        conv = _conv_with([make_image_placeholder(ref, "user_upload")])
        report = LossReport()
        outcome = resolve_media(
            conv, provider, tmp_path, "provider/project/year", "images", report,
            circuit=circuit,
        )
        return outcome, report, conv

    def test_server_error_is_retryable(self, tmp_path):
        outcome, report, conv = self._run("Signed URL: HTTP 500 after 5 retries", tmp_path)
        assert outcome.failed == 1
        assert outcome.retryable_refs == ["sediment://file_x"]
        assert report.media_failed["download-error"] == 1
        # Placeholder survives, so the ref is still in the Markdown.
        assert "local_path" not in conv["messages"][0]["blocks"][0]

    def test_service_unavailable_is_retryable(self, tmp_path):
        outcome, _, _ = self._run("Signed URL: HTTP 503 after 5 retries", tmp_path)
        assert outcome.retryable_refs == ["sediment://file_x"]

    def test_timeout_is_retryable(self, tmp_path):
        outcome, _, _ = self._run(
            "Network error after 5 retries: Operation timed out", tmp_path
        )
        assert outcome.retryable_refs == ["sediment://file_x"]

    def test_deleted_asset_is_not_retryable(self, tmp_path):
        outcome, report, _ = self._run(
            "Asset no longer exists — HTTP 404 'File not found'", tmp_path
        )
        assert outcome.failed == 1
        assert outcome.retryable_refs == []
        assert report.media_failed["expired-or-missing"] == 1

    def test_forbidden_asset_is_not_retryable(self, tmp_path):
        outcome, report, _ = self._run("HTTP 403 — detail: Forbidden", tmp_path)
        assert outcome.retryable_refs == []
        assert report.media_failed["forbidden"] == 1

    def test_success_reports_nothing_retryable(self, tmp_path):
        from src.media import resolve_media

        ref = "sediment://file_x"
        provider = _FakeProvider(assets={ref: (b"png", "image/png", None)})
        conv = _conv_with([make_image_placeholder(ref, "user_upload")])
        outcome = resolve_media(
            conv, provider, tmp_path, "provider/project/year", "images", LossReport()
        )
        assert (outcome.downloaded, outcome.failed, outcome.retryable_refs) == (1, 0, [])


class TestMediaCircuit:
    def test_trips_after_consecutive_transient_failures(self):
        from src.media import MediaCircuit

        c = MediaCircuit(threshold=3)
        for _ in range(2):
            c.record_failure("download-error")
        assert not c.tripped
        c.record_failure("download-error")
        assert c.tripped

    def test_success_resets_the_streak(self):
        from src.media import MediaCircuit

        c = MediaCircuit(threshold=3)
        c.record_failure("download-error")
        c.record_failure("download-error")
        c.record_success()
        c.record_failure("download-error")
        assert not c.tripped, "scattered failures must not trip the breaker"

    def test_permanent_failures_never_trip_it(self):
        """A deleted upload says nothing about the provider's health."""
        from src.media import MediaCircuit

        c = MediaCircuit(threshold=2)
        for _ in range(10):
            c.record_failure("expired-or-missing")
            c.record_failure("forbidden")
        assert not c.tripped

    def test_open_circuit_skips_downloads_but_keeps_refs(self, tmp_path):
        from src.media import CIRCUIT_OPEN_BUCKET, MediaCircuit, resolve_media

        ref = "sediment://file_y"
        provider = _FakeProvider(assets={ref: (b"png", "image/png", None)})
        conv = _conv_with([make_image_placeholder(ref, "user_upload")])
        circuit = MediaCircuit(threshold=1)
        circuit.record_failure("download-error")
        assert circuit.tripped

        report = LossReport()
        outcome = resolve_media(
            conv, provider, tmp_path, "provider/project/year", "images", report,
            circuit=circuit,
        )
        assert provider.calls == [], "must not touch the network while open"
        assert outcome.downloaded == 0
        assert outcome.retryable_refs == [ref], "skipped assets must not be abandoned"
        assert report.media_failed[CIRCUIT_OPEN_BUCKET] == 1
        assert circuit.skipped == 1

    def test_already_downloaded_asset_served_while_open(self, tmp_path):
        """An open circuit must not stop on-disk media being linked."""
        from src.media import MediaCircuit, resolve_media

        ref = "sediment://file_z"
        conv = _conv_with([make_image_placeholder(ref, "user_upload")])
        # Derive the directory the way the exporter does rather than
        # hardcoding it — the layout is "<provider>/<project>.<year>/".
        from src.utils import build_export_path, generate_filename

        media_dir = build_export_path(
            tmp_path,
            "chatgpt",
            None,
            "2026-05-20T00:00:00+00:00",
            generate_filename("Has Media", "conv-1", "2026-05-20T00:00:00+00:00"),
            "provider/project/year",
        ).parent / "media"
        media_dir.mkdir(parents=True)
        (media_dir / "file_z.png").write_bytes(b"already here")

        circuit = MediaCircuit(threshold=1)
        circuit.record_failure("download-error")
        provider = _FakeProvider()
        resolve_media(
            conv, provider, tmp_path, "provider/project/year", "images", LossReport(),
            circuit=circuit,
        )
        assert provider.calls == []
        assert conv["messages"][0]["blocks"][0]["local_path"] == "media/file_z.png"