mcp_server/tests/test_catalog.py

from __future__ import annotations

import json

import pytest
from pwc_cli.cli import UsageError
from pwc_cli.transport import Response, ResponseError
from pwc_mcp.catalog import (
    CatalogClient,
    PublicPaperMiss,
    cache_ttl,
    near_title,
    normalize_title,
)


class StubTransport:
    def __init__(self, routes):
        self.routes = routes
        self.calls = []

    def get(self, path, params=None):
        self.calls.append((path, params))
        value = self.routes[path]
        if isinstance(value, str):
            return Response(value.encode(), {"content-type": "text/markdown"})
        return Response(
            json.dumps(value).encode(), {"content-type": "application/json"}
        )


def test_catalog_normalizes_paper_urls_and_caches_identical_reads():
    transport = StubTransport(
        {
            "papers/1706.03762": {"id": "755", "title": "Attention Is All You Need"},
        }
    )
    catalog = CatalogClient(transport=transport)

    first = catalog.query(
        ("paper", "info"),
        {"paper": "https://arxiv.org/pdf/1706.03762v7.pdf", "include_resources": True},
    )
    second = catalog.query(
        ("paper", "info"),
        {
            "paper": "https://huggingface.co/papers/1706.03762",
            "include_resources": True,
        },
    )

    assert first == second
    assert transport.calls == [("papers/1706.03762", {"include_resources": True})]


def test_catalog_reads_one_version_bound_markdown_chunk_per_call():
    version = "a" * 64

    class ChunkTransport:
        def __init__(self):
            self.calls = []

        def get(self, path, params=None):
            self.calls.append((path, params))
            if params["offset"] == 0:
                return Response(
                    "α first".encode(),
                    {
                        "x-pwc-truncated": "1",
                        "x-pwc-content-version": version,
                        "x-pwc-next-offset": "8",
                    },
                )
            return Response(
                b" terminal",
                {
                    "x-pwc-truncated": "0",
                    "x-pwc-content-version": version,
                },
            )

    transport = ChunkTransport()
    catalog = CatalogClient(transport=transport)
    first = catalog.read_paper_chunk("1706.03762", limit=65_536)
    second = catalog.read_paper_chunk(
        "1706.03762",
        offset=first.next_offset,
        content_version=first.content_version,
        limit=65_536,
        resolved=True,
    )

    assert first.markdown == "α first"
    assert first.next_offset == 8
    assert second.markdown == " terminal"
    assert second.next_offset is None
    assert transport.calls == [
        (
            "research/papers/1706.03762/read",
            {"offset": 0, "limit": 65_536, "content_version": None},
        ),
        (
            "research/papers/1706.03762/read",
            {"offset": 8, "limit": 65_536, "content_version": version},
        ),
    ]


def test_catalog_caches_versioned_chunks_and_rejects_nonadvancing_offsets():
    version = "b" * 64

    class ChunkTransport:
        def __init__(self):
            self.calls = 0

        def get(self, _path, _params=None):
            self.calls += 1
            return Response(
                b"data",
                {
                    "x-pwc-truncated": "1",
                    "x-pwc-content-version": version,
                    "x-pwc-next-offset": "4",
                },
            )

    transport = ChunkTransport()
    catalog = CatalogClient(transport=transport)
    with pytest.raises(ResponseError, match="did not advance"):
        catalog.read_paper_chunk(
            "1706.03762",
            offset=4,
            content_version=version,
            resolved=True,
        )
    assert transport.calls == 1


def test_catalog_versioned_chunk_cache_avoids_repeated_upstream_reads():
    version = "c" * 64

    class ChunkTransport:
        def __init__(self):
            self.calls = 0

        def get(self, _path, _params=None):
            self.calls += 1
            return Response(
                b"done",
                {
                    "x-pwc-truncated": "0",
                    "x-pwc-content-version": version,
                },
            )

    transport = ChunkTransport()
    catalog = CatalogClient(transport=transport)
    first = catalog.read_paper_chunk(
        "1706.03762",
        offset=8,
        content_version=version,
        resolved=True,
    )
    second = catalog.read_paper_chunk(
        "1706.03762",
        offset=8,
        content_version=version,
        resolved=True,
    )

    assert first == second
    assert transport.calls == 1


def test_catalog_resolves_exact_titles_and_rejects_ambiguous_titles():
    transport = StubTransport(
        {
            "papers/search": {
                "results": [
                    {"id": "1", "arxiv_id": "1111.11111", "title": "Same Title"},
                    {"id": "2", "arxiv_id": "2222.22222", "title": "Same Title"},
                ]
            }
        }
    )
    catalog = CatalogClient(transport=transport)

    with pytest.raises(ResponseError, match="ambiguous"):
        catalog.query(("paper", "lineage", "list"), {"paper": "Same Title"})


def test_catalog_resolves_pwc_urls_and_dotted_legacy_arxiv_ids():
    transport = StubTransport(
        {
            "papers/search": {
                "results": [
                    {
                        "id": "755",
                        "arxiv_id": "1706.03762",
                        "title": "Attention Is All You Need",
                    }
                ]
            },
            "papers/1706.03762": {"id": "755"},
            "papers/math.GT%2F0309136": {"id": "900"},
        }
    )
    catalog = CatalogClient(transport=transport)

    assert catalog.query(
        ("paper", "info"),
        {"paper": "https://paperswithcode.co/paper/attention-is-all-you-need"},
    ) == {"id": "755"}
    assert catalog.query(("paper", "info"), {"paper": "math.GT/0309136"}) == {
        "id": "900"
    }


def test_catalog_checks_later_search_pages_for_ambiguous_exact_titles():
    class PaginatedTransport:
        def get(self, path, params=None):
            assert path == "papers/search"
            page = params["page"]
            payload = {
                "results": [
                    {
                        "id": str(page),
                        "arxiv_id": f"1111.1111{page}",
                        "title": "Same Title",
                    }
                ],
                "next_page": 2 if page == 1 else None,
            }
            return Response(json.dumps(payload).encode(), {})

    catalog = CatalogClient(transport=PaginatedTransport())

    with pytest.raises(ResponseError, match="ambiguous"):
        catalog.query(("paper", "lineage", "list"), {"paper": "Same Title"})


def test_catalog_fails_closed_when_paper_filters_are_not_confirmed():
    transport = StubTransport({"papers/": {"results": [], "applied_filters": {}}})
    catalog = CatalogClient(transport=transport)

    with pytest.raises(ResponseError, match="did not confirm"):
        catalog.query(("paper", "list"), {"task": "image-classification"})


def test_catalog_query_runs_cli_commands_with_per_path_cache_lifetimes():
    transport = StubTransport(
        {
            "datasets/": {"results": [{"id": "72", "name": "ImageNet-1k"}]},
            "datasets/72/evaluations/": {
                "count": 1,
                "results": [{"id": "1", "model_name": "A", "metrics": {"top1": 80}}],
            },
        }
    )
    catalog = CatalogClient(transport=transport)

    first = catalog.query(("benchmark",), {"name": "ImageNet-1k", "limit": 5})
    second = catalog.query(("benchmark",), {"name": "imagenet-1k", "limit": 5})

    assert first == second
    assert first["results"][0]["model_name"] == "A"
    assert first["count"] == 1
    # The second lookup differs only in the search text; the identical
    # evaluation request is served from the shared cache.
    assert [path for path, _params in transport.calls] == [
        "datasets/",
        "datasets/72/evaluations/",
        "datasets/",
    ]
    assert cache_ttl("papers/search") == 60
    assert cache_ttl("tasks/1/page") == 600
    assert cache_ttl("datasets/72/evaluations/") == 300


def test_catalog_query_reports_cli_usage_errors_without_upstream_calls():
    transport = StubTransport({})
    catalog = CatalogClient(transport=transport)

    with pytest.raises(UsageError, match="--limit must be 1-100"):
        catalog.query(("search",), {"query": "x", "limit": 0})
    with pytest.raises(UsageError, match="not a read-only"):
        catalog.query(("skills", "add"), {})
    assert transport.calls == []


def test_catalog_canonicalises_numeric_ids_of_arxiv_papers():
    transport = StubTransport(
        {
            "papers/86122": {"id": "86122", "arxiv_id": "2505.18132v1"},
            "papers/2505.18132": {"id": "86122", "title": "BiggerGait"},
            "papers/4242": {"id": "4242", "arxiv_id": None, "source": "external"},
        }
    )
    catalog = CatalogClient(transport=transport)

    # List tools hand out numeric IDs for arXiv papers too; the Markdown route
    # only knows arXiv papers by arXiv ID, so numeric references canonicalise.
    assert catalog.resolve_paper("86122") == "2505.18132"
    assert catalog.query(("paper", "info"), {"paper": "86122"}) == {
        "id": "86122",
        "title": "BiggerGait",
    }
    # External papers have no arXiv ID and keep their numeric reference.
    assert catalog.resolve_paper("4242") == "4242"
    assert transport.calls[0] == ("papers/86122", {})


def test_catalog_rejects_unsupported_urls_without_searching():
    transport = StubTransport({})
    catalog = CatalogClient(transport=transport)

    for url in (
        "https://doi.org/10.1007/s10479-024-06277-x",
        "https://aclanthology.org/2026.acl-long.200/",
    ):
        with pytest.raises(ResponseError, match="Paper URL not supported"):
            catalog.resolve_paper(url)
    assert transport.calls == []


def test_catalog_distinguishes_missing_papers_from_missing_markdown():
    from pwc_cli.transport import HTTPStatusError

    class MissingTransport:
        def __init__(self, paper_exists):
            self.paper_exists = paper_exists

        def get(self, path, params=None):
            if path.endswith("/read"):
                raise HTTPStatusError(404, "Paper Markdown was not found")
            if self.paper_exists:
                return Response(b'{"id": "1"}', {})
            raise HTTPStatusError(404, "No paper found")

    with pytest.raises(ResponseError, match="Paper not found: 2308.10195"):
        CatalogClient(transport=MissingTransport(False)).read_paper_chunk(
            "2308.10195", resolved=True
        )
    with pytest.raises(HTTPStatusError) as missing_markdown:
        CatalogClient(transport=MissingTransport(True)).read_paper_chunk(
            "2505.18132", resolved=True
        )
    assert missing_markdown.value.status == 404


def test_catalog_reports_titles_missing_from_an_empty_search_as_not_found():
    # The public API answers an unknown title with an empty result list; that
    # must surface as "title not found", not as a malformed response.
    transport = StubTransport(
        {
            "papers/search": {
                "next_page": None,
                "previous_page": None,
                "results": [],
                "applied_filters": {},
            }
        }
    )
    catalog = CatalogClient(transport=transport)

    with pytest.raises(ResponseError, match="Paper title not found: Dropout"):
        catalog.resolve_paper("Dropout: A Simple Way to Prevent Overfitting")
    # One phrase query, then one plain keyword query as the fallback.
    assert [params["q"] for _path, params in transport.calls] == [
        '"Dropout: A Simple Way to Prevent Overfitting"',
        "Dropout: A Simple Way to Prevent Overfitting",
    ]


EXACT_TITLE_ROW = {"id": "755", "arxiv_id": "1706.03762", "title": "Attention Is All You Need"}


class TitleSearchTransport:
    """Keyword search whose plain results never end; the phrase query is small."""

    def __init__(self, *, phrase_results, plain_has_exact):
        self.phrase_results = phrase_results
        self.plain_has_exact = plain_has_exact
        self.calls = []

    def get(self, path, params=None):
        params = dict(params or {})
        self.calls.append((path, params))
        assert path == "papers/search"
        page = params["page"]
        if params["q"].startswith('"'):
            body = {"results": self.phrase_results, "next_page": None}
        else:
            rows = [
                {"id": str(page * 100 + i), "arxiv_id": f"2{page:03d}.{i:05d}", "title": f"Attention {i}"}
                for i in range(100)
            ]
            if self.plain_has_exact and page == 1:
                rows[0] = EXACT_TITLE_ROW
            body = {"results": rows, "next_page": page + 1}
        return Response(json.dumps(body).encode(), {"content-type": "application/json"})


def test_catalog_resolves_common_word_titles_through_one_phrase_query():
    # "Attention Is All You Need" has over a thousand keyword hits; paging
    # through them used to end in "Too many results" although the exact title
    # was the first result.
    transport = TitleSearchTransport(
        phrase_results=[
            EXACT_TITLE_ROW,
            {"id": "9", "arxiv_id": "2010.13154", "title": "Attention is All You Need in Speech Separation"},
        ],
        plain_has_exact=True,
    )
    catalog = CatalogClient(transport=transport)

    assert catalog.resolve_paper("Attention Is All You Need") == "1706.03762"
    assert [params["q"] for _path, params in transport.calls] == ['"Attention Is All You Need"']


def test_catalog_keeps_exact_matches_found_before_the_page_budget_runs_out():
    transport = TitleSearchTransport(phrase_results=[], plain_has_exact=True)
    catalog = CatalogClient(transport=transport)

    assert catalog.resolve_paper("Attention Is All You Need") == "1706.03762"
    assert len(transport.calls) == 11

    inconclusive = CatalogClient(
        transport=TitleSearchTransport(phrase_results=[], plain_has_exact=False)
    )
    with pytest.raises(ResponseError, match="Too many results"):
        inconclusive.resolve_paper("Attention Is All You Need")


def test_catalog_resolves_a_unique_short_name_before_the_title_colon():
    # Production chat: get_paper_info("MuseTalk") failed although the catalog
    # holds "MuseTalk: Real-Time High-Fidelity Video Dubbing ..." (2410.10122).
    musetalk = {
        "id": "90",
        "arxiv_id": "2410.10122",
        "title": "MuseTalk: Real-Time High-Fidelity Video Dubbing via Spatio-Temporal Sampling",
    }
    citing = {"id": "91", "arxiv_id": "2501.00001", "title": "Beyond MuseTalk: Faster Dubbing"}
    transport = TitleSearchTransport(phrase_results=[musetalk, citing], plain_has_exact=False)
    catalog = CatalogClient(transport=transport)

    assert catalog.resolve_paper("MuseTalk") == "2410.10122"
    assert catalog.resolve_paper("musetalk") == "2410.10122"

    twins = [musetalk, {**musetalk, "id": "92", "arxiv_id": "2502.00002"}]
    ambiguous = CatalogClient(
        transport=TitleSearchTransport(phrase_results=twins, plain_has_exact=False)
    )
    with pytest.raises(ResponseError, match="Paper title is ambiguous: MuseTalk"):
        ambiguous.resolve_paper("MuseTalk")


def test_catalog_resolves_a_unique_near_match_of_a_full_title():
    # Production chat: get_paper_info("BEIR: A Heterogeneous Benchmark ...")
    # failed although the catalog holds the arXiv spelling "Heterogenous"
    # (2104.08663), which keyword search returns first.
    beir = {
        "id": "5505",
        "arxiv_id": "2104.08663",
        "title": "BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models",
    }
    beir_nl = {
        "id": "93",
        "arxiv_id": "2412.08329",
        "title": "BEIR-NL: Zero-shot Information Retrieval Benchmark for the Dutch Language",
    }
    catalog = CatalogClient(
        transport=TitleSearchTransport(phrase_results=[beir, beir_nl], plain_has_exact=False)
    )

    assert catalog.resolve_paper(
        "BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models"
    ) == "2104.08663"

    twins = [beir, {**beir, "id": "94", "arxiv_id": "2502.00003"}]
    ambiguous = CatalogClient(
        transport=TitleSearchTransport(phrase_results=twins, plain_has_exact=False)
    )
    with pytest.raises(ResponseError, match="Paper title is ambiguous"):
        ambiguous.resolve_paper(
            "BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models"
        )


def test_near_title_rejects_short_names_and_other_editions():
    assert near_title(
        normalize_title("BEIR: A Heterogenous Benchmark for Zero-shot Evaluation"),
        normalize_title("BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation"),
    )
    assert not near_title(normalize_title("MuseTalks"), normalize_title("MuseTalk"))
    assert not near_title(
        normalize_title("The Llama 3 Herd of Language Foundation Models"),
        normalize_title("The Llama 4 Herd of Language Foundation Models"),
    )
    assert not near_title(
        normalize_title("Deep Residual Learning for Image Recognition"),
        normalize_title("Deep Residual Learning for Video Recognition"),
    )


def test_catalog_reads_public_metadata_and_caches_lookups_including_misses():
    from pwc_cli.transport import HTTPStatusError

    class PublicTransport:
        def __init__(self):
            self.calls = []

        def get(self, path, params=None):
            self.calls.append((path, params))
            if path == "papers/2610.18869/public-metadata":
                return Response(
                    json.dumps({"in_catalog": False, "title": "Found"}).encode(), {}
                )
            if path == "papers/public-metadata/search":
                title = params["title"]
                if title == "Down Title":
                    raise HTTPStatusError(503, "unavailable")
                if title == "BERT":
                    raise HTTPStatusError(400, "title too short")
                if title == "Some Title":
                    return Response(
                        json.dumps({"in_catalog": False, "title": "Some title"}).encode(),
                        {},
                    )
                raise HTTPStatusError(404, "no exact title")
            if path.startswith("papers/Down/"):
                raise HTTPStatusError(503, "unavailable")
            if "%20" in path or path == "papers/BERT/public-metadata":
                raise HTTPStatusError(400, "titles are not identifiers")
            raise HTTPStatusError(404, "Paper not found")

    transport = PublicTransport()
    catalog = CatalogClient(transport=transport)

    assert catalog.public_paper(" 2610.18869 ")["title"] == "Found"
    assert catalog.public_paper("2610.18869")["title"] == "Found"
    assert catalog.public_paper("10.1/x") == PublicPaperMiss("identifier")
    assert catalog.public_paper("10.1/x") == PublicPaperMiss("identifier")
    # Titles the identifier route rejects are searched as exact titles.
    assert catalog.public_paper("Some Title") == {
        "in_catalog": False,
        "title": "Some title",
        "requested_title": "Some Title",
    }
    assert catalog.public_paper("Some Title")["title"] == "Some title"
    assert catalog.public_paper("Unknown Title") == PublicPaperMiss("title")
    assert catalog.public_paper("Unknown Title") == PublicPaperMiss("title")
    assert catalog.public_paper("BERT") == PublicPaperMiss("unsupported")
    # Upstream failures are best effort and not cached as misses.
    assert catalog.public_paper("Down") == PublicPaperMiss("unavailable")
    assert catalog.public_paper("Down") == PublicPaperMiss("unavailable")
    assert catalog.public_paper("Down Title") == PublicPaperMiss("unavailable")
    assert catalog.public_paper("Down Title") == PublicPaperMiss("unavailable")
    search = "papers/public-metadata/search"
    assert transport.calls == [
        ("papers/2610.18869/public-metadata", None),
        ("papers/10.1/x/public-metadata", None),
        ("papers/Some%20Title/public-metadata", None),
        (search, {"title": "Some Title"}),
        ("papers/Unknown%20Title/public-metadata", None),
        (search, {"title": "Unknown Title"}),
        ("papers/BERT/public-metadata", None),
        (search, {"title": "BERT"}),
        ("papers/Down/public-metadata", None),
        ("papers/Down/public-metadata", None),
        ("papers/Down%20Title/public-metadata", None),
        (search, {"title": "Down Title"}),
        ("papers/Down%20Title/public-metadata", None),
        (search, {"title": "Down Title"}),
    ]