codex-py/tests/sources/test_arxiv.py

"""Tests for codex.sources.arxiv."""

from __future__ import annotations

import io
import tarfile

import httpx
import pytest

from codex.sources import arxiv

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_tar_gz(files: dict[str, bytes]) -> bytes:
    """Build an in-memory .tar.gz archive from a dict of {name: content}."""
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tf:
        for name, content in files.items():
            info = tarfile.TarInfo(name=name)
            info.size = len(content)
            tf.addfile(info, io.BytesIO(content))
    return buf.getvalue()


# ---------------------------------------------------------------------------
# fetch_source — success: in-memory tar.gz with one .tex
# ---------------------------------------------------------------------------


def test_fetch_source_returns_latex(monkeypatch: pytest.MonkeyPatch) -> None:
    """fetch_source extracts and returns the LaTeX content from a .tar.gz."""
    latex_content = b"\\documentclass{article}\n\\begin{document}\nHello world\n\\end{document}"
    tar_bytes = _make_tar_gz({"main.tex": latex_content})

    def mock_get(
        url: str,
        *,
        timeout: int = 60,
        follow_redirects: bool = True,
    ) -> httpx.Response:
        return httpx.Response(200, content=tar_bytes)

    monkeypatch.setattr(httpx, "get", mock_get)

    result = arxiv.fetch_source("2301.07041")

    assert result is not None
    assert "\\documentclass" in result
    assert "Hello world" in result


def test_fetch_source_flattens_multifile(monkeypatch: pytest.MonkeyPatch) -> None:
    """Multi-file projects: \\input/\\include directives are inlined so the body
    from the included files is returned, not just the primary skeleton (R-C)."""
    main = (
        b"\\documentclass{article}\\begin{document}"
        b"\\input{sections/intro}\\include{proof}\\end{document}"
    )
    tar_bytes = _make_tar_gz(
        {
            "main.tex": main,
            "sections/intro.tex": b"INTRODUCTION BODY TEXT",
            "proof.tex": b"PROOF BODY TEXT",
        }
    )

    def mock_get(
        url: str,
        *,
        timeout: int = 60,
        follow_redirects: bool = True,
    ) -> httpx.Response:
        return httpx.Response(200, content=tar_bytes)

    monkeypatch.setattr(httpx, "get", mock_get)

    result = arxiv.fetch_source("2301.07041")

    assert result is not None
    assert "INTRODUCTION BODY TEXT" in result  # \input was inlined
    assert "PROOF BODY TEXT" in result  # \include was inlined
    assert "\\input" not in result and "\\include" not in result


# ---------------------------------------------------------------------------
# fetch_source — 404 → None
# ---------------------------------------------------------------------------


def test_fetch_source_404_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
    """A 404 response should return None."""

    def mock_get(
        url: str,
        *,
        timeout: int = 60,
        follow_redirects: bool = True,
    ) -> httpx.Response:
        return httpx.Response(404)

    monkeypatch.setattr(httpx, "get", mock_get)

    result = arxiv.fetch_source("nonexistent-id")
    assert result is None


# ---------------------------------------------------------------------------
# fetch_source — no .tex in archive → None
# ---------------------------------------------------------------------------


def test_fetch_source_no_tex_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
    """If the archive contains no .tex files, return None."""
    tar_bytes = _make_tar_gz({"README.md": b"# Paper\nNo LaTeX here."})

    def mock_get(
        url: str,
        *,
        timeout: int = 60,
        follow_redirects: bool = True,
    ) -> httpx.Response:
        return httpx.Response(200, content=tar_bytes)

    monkeypatch.setattr(httpx, "get", mock_get)

    result = arxiv.fetch_source("2301.07041")
    assert result is None


# ---------------------------------------------------------------------------
# fetch_source — fallback to largest .tex when no \documentclass
# ---------------------------------------------------------------------------


def test_fetch_source_fallback_to_largest_tex(monkeypatch: pytest.MonkeyPatch) -> None:
    """When no file has \\documentclass, the largest .tex is returned."""
    small_tex = b"% small file\n\\section{Intro}"
    large_tex = b"% large file\n" + b"x" * 500

    tar_bytes = _make_tar_gz({"small.tex": small_tex, "large.tex": large_tex})

    def mock_get(
        url: str,
        *,
        timeout: int = 60,
        follow_redirects: bool = True,
    ) -> httpx.Response:
        return httpx.Response(200, content=tar_bytes)

    monkeypatch.setattr(httpx, "get", mock_get)

    result = arxiv.fetch_source("2301.07041")

    assert result is not None
    assert "large file" in result


# ---------------------------------------------------------------------------
# fetch_pdf_url — pure computation, no mock needed
# ---------------------------------------------------------------------------


def test_fetch_pdf_url_pure_computation() -> None:
    """fetch_pdf_url returns the correct URL without making any HTTP request."""
    url = arxiv.fetch_pdf_url("2301.07041")
    assert url == "https://arxiv.org/pdf/2301.07041.pdf"

    url2 = arxiv.fetch_pdf_url("1234.56789")
    assert url2 == "https://arxiv.org/pdf/1234.56789.pdf"


# ---------------------------------------------------------------------------
# fetch_metadata — parse the Atom feed into a Paper (DQ-2)
# ---------------------------------------------------------------------------

_ATOM_FEED = """<?xml version="1.0" encoding="UTF-8"?>
<feed xmlns="http://www.w3.org/2005/Atom">
  <entry>
    <title>A discrete version of Liouville's theorem on conformal maps</title>
    <summary>Liouville's theorem says that in dimension greater than two,
    all conformal maps are Moebius transformations.</summary>
    <published>2019-11-03T00:00:00Z</published>
    <author><name>Ulrich Pinkall</name></author>
    <author><name>Boris Springborn</name></author>
  </entry>
</feed>"""

_ATOM_EMPTY = """<?xml version="1.0" encoding="UTF-8"?>
<feed xmlns="http://www.w3.org/2005/Atom"><title>ArXiv Query</title></feed>"""


def test_fetch_metadata_parses_entry(monkeypatch: pytest.MonkeyPatch) -> None:
    """A valid Atom feed yields a populated Paper."""

    def mock_get(url: str, **kwargs: object) -> httpx.Response:
        return httpx.Response(200, text=_ATOM_FEED, request=httpx.Request("GET", url))

    monkeypatch.setattr(httpx, "get", mock_get)

    paper = arxiv.fetch_metadata("1911.00966")
    assert paper is not None
    assert paper.id == "1911.00966"
    assert paper.title.startswith("A discrete version of Liouville")
    assert paper.authors == ["Ulrich Pinkall", "Boris Springborn"]
    assert paper.year == 2019
    assert paper.abstract and "conformal maps" in paper.abstract
    assert paper.openalex_id is None


def test_fetch_metadata_no_entry_returns_none(monkeypatch: pytest.MonkeyPatch) -> None:
    """A feed with no <entry> (id not found) returns None."""

    def mock_get(url: str, **kwargs: object) -> httpx.Response:
        return httpx.Response(200, text=_ATOM_EMPTY, request=httpx.Request("GET", url))

    monkeypatch.setattr(httpx, "get", mock_get)
    assert arxiv.fetch_metadata("0000.00000") is None


def test_fetch_metadata_strips_arxiv_prefix(monkeypatch: pytest.MonkeyPatch) -> None:
    """An ``arXiv:`` prefix is stripped so Paper.id is the bare id."""

    def mock_get(url: str, **kwargs: object) -> httpx.Response:
        return httpx.Response(200, text=_ATOM_FEED, request=httpx.Request("GET", url))

    monkeypatch.setattr(httpx, "get", mock_get)
    paper = arxiv.fetch_metadata("arXiv:1911.00966")
    assert paper is not None
    assert paper.id == "1911.00966"