codex-py/tests/sources/test_openalex.py

"""Tests for codex.sources.openalex."""

from __future__ import annotations

import httpx
import pytest

from codex.models import Citation, Paper
from codex.sources import openalex
from codex.sources.openalex import _resolve_id

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

_SAMPLE_WORK = {
    "id": "https://openalex.org/W2741809807",
    "doi": "10.1145/3592430",
    "title": "Conformal Prediction: A Review",
    "publication_year": 2023,
    "authorships": [
        {"author": {"display_name": "Alice Smith"}},
        {"author": {"display_name": "Bob Jones"}},
    ],
    "abstract_inverted_index": {
        "Conformal": [0],
        "prediction": [1],
        "is": [2],
        "great": [3],
    },
    "referenced_works": [
        "https://openalex.org/W1",
        "https://openalex.org/W2",
    ],
}


# ---------------------------------------------------------------------------
# _resolve_id
# ---------------------------------------------------------------------------


def test_resolve_doi() -> None:
    assert _resolve_id("10.1145/3592430") == "doi:10.1145/3592430"


def test_resolve_arxiv() -> None:
    # arXiv IDs resolve to their arXiv DOI (the arxiv: path 404s on OpenAlex).
    assert _resolve_id("2301.07041") == "doi:10.48550/arXiv.2301.07041"


def test_resolve_arxiv_legacy() -> None:
    # Legacy category/number IDs (pre-2007) resolve the same way.
    assert _resolve_id("math/0603097") == "doi:10.48550/arXiv.math/0603097"


def test_resolve_arxiv_prefix_stripped() -> None:
    # An explicit arxiv: prefix is stripped before building the arXiv DOI.
    assert _resolve_id("arxiv:2301.07041") == "doi:10.48550/arXiv.2301.07041"


def test_resolve_w_id_unchanged() -> None:
    assert _resolve_id("W2741809807") == "W2741809807"


def test_resolve_prefixed_unchanged() -> None:
    assert _resolve_id("doi:10.1145/x") == "doi:10.1145/x"


# ---------------------------------------------------------------------------
# fetch_paper — success
# ---------------------------------------------------------------------------


def test_fetch_paper_success(monkeypatch: pytest.MonkeyPatch) -> None:
    """Successful fetch maps all fields to Paper correctly."""

    def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response:
        assert "doi:10.1145/3592430" in url
        return httpx.Response(200, json=_SAMPLE_WORK)

    monkeypatch.setattr(openalex, "_get", mock_get)

    paper = openalex.fetch_paper("10.1145/3592430")

    assert paper is not None
    assert isinstance(paper, Paper)
    assert paper.title == "Conformal Prediction: A Review"
    assert paper.year == 2023
    assert paper.authors == ["Alice Smith", "Bob Jones"]
    assert paper.abstract == "Conformal prediction is great"
    assert paper.openalex_id == "https://openalex.org/W2741809807"


def test_fetch_paper_arxiv_id_prefixed(monkeypatch: pytest.MonkeyPatch) -> None:
    """Bare arXiv IDs are resolved to their arXiv DOI in the URL."""
    called_url: list[str] = []

    def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response:
        called_url.append(url)
        return httpx.Response(200, json=_SAMPLE_WORK)

    monkeypatch.setattr(openalex, "_get", mock_get)
    openalex.fetch_paper("2301.07041")

    assert called_url and "doi:10.48550/arXiv.2301.07041" in called_url[0]


# ---------------------------------------------------------------------------
# fetch_paper — 404 → None
# ---------------------------------------------------------------------------


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

    def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response:
        raise httpx.HTTPStatusError(
            "Not Found",
            request=httpx.Request("GET", url),
            response=httpx.Response(404, request=httpx.Request("GET", url)),
        )

    monkeypatch.setattr(openalex, "_get", mock_get)

    result = openalex.fetch_paper("nonexistent-id")
    assert result is None


# ---------------------------------------------------------------------------
# fetch_paper — 429 retry then success
# ---------------------------------------------------------------------------


def test_fetch_paper_retry_on_429(monkeypatch: pytest.MonkeyPatch) -> None:
    """On 429×3 then 200, returns 1 Paper after 3 failed attempts."""
    attempt = [0]

    def controlled_get(url: str, params: dict[str, str] | None = None) -> httpx.Response:
        attempt[0] += 1
        if attempt[0] <= 3:
            resp = httpx.Response(429, request=httpx.Request("GET", url))
            raise httpx.HTTPStatusError(
                "Too Many Requests",
                request=httpx.Request("GET", url),
                response=resp,
            )
        return httpx.Response(200, json=_SAMPLE_WORK)

    monkeypatch.setattr(openalex, "_get", controlled_get)

    result: Paper | None = None
    errors = 0
    for _ in range(5):
        try:
            result = openalex.fetch_paper("10.1145/3592430")
            break
        except httpx.HTTPStatusError:
            errors += 1

    assert result is not None
    assert isinstance(result, Paper)
    assert errors == 3


# ---------------------------------------------------------------------------
# fetch_citations — uses referenced_works field
# ---------------------------------------------------------------------------


def test_fetch_citations(monkeypatch: pytest.MonkeyPatch) -> None:
    """fetch_citations reads referenced_works from the work object."""

    def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response:
        return httpx.Response(200, json=_SAMPLE_WORK)

    monkeypatch.setattr(openalex, "_get", mock_get)

    oa_id = "https://openalex.org/W2741809807"
    citations = openalex.fetch_citations(oa_id)

    assert len(citations) == 2
    assert all(isinstance(c, Citation) for c in citations)
    assert citations[0].citing_id == oa_id
    assert citations[0].cited_id == "https://openalex.org/W1"
    assert citations[1].cited_id == "https://openalex.org/W2"


def test_fetch_citations_404_returns_empty(monkeypatch: pytest.MonkeyPatch) -> None:
    """404 on the work fetch returns an empty list."""

    def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response:
        raise httpx.HTTPStatusError(
            "Not Found",
            request=httpx.Request("GET", url),
            response=httpx.Response(404, request=httpx.Request("GET", url)),
        )

    monkeypatch.setattr(openalex, "_get", mock_get)
    assert openalex.fetch_citations("W999") == []