user2595/codex-py
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
d64edfc3249f4d4feb3733983df120ae34ee09c0
codex-py/tests/ingest/test_ingest.py
Tarik Moussa e1c0a98262 fix(ingest): make both upsert UNIQUE gaps visible (audit C-15 follow-up) 
Per review: the two latent upsert conflicts should be observable even where not
auto-fixed.
- bibkey collision (auto-suffixed) now logs a warning so a paper landing as
  'BobenkoLutz2023a' instead of '…2023' isn't silent.
- openalex_id collision (two papers claiming the same OpenAlex work — a real data
  conflict, not auto-renamable) now raises a clear ValueError naming the
  openalex_id and the offending paper, instead of the raw psycopg UniqueViolation.

Regression test added for the openalex_id error path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-15 22:25:30 +02:00

610 lines
24 KiB
Python
Raw Blame History

"""Tests for codex.ingest — end-to-end idempotent ingest pipeline.
All external calls (DB, OpenAlex, S2, embed, parsing) are mocked so
the suite runs offline in milliseconds.
"""
from __future__ import annotations
from contextlib import contextmanager
from typing import Any
from unittest.mock import MagicMock, patch
import numpy as np
import psycopg
import pytest
from codex.ingest import IngestResult, _make_bibkey, ingest_paper
from codex.models import Citation, FigureChunk, FormulaChunk, Paper
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_conn_cm(mock_conn: MagicMock) -> Any:
"""Return a context-manager factory that yields *mock_conn*."""
@contextmanager # type: ignore[arg-type]
def _cm() -> Any:
yield mock_conn
return _cm
def _make_paper(
paper_id: str = "2301.07041",
openalex_id: str | None = "W123",
abstract: str | None = "Test abstract.",
) -> Paper:
return Paper(
id=paper_id,
title="A Test Paper",
openalex_id=openalex_id,
authors=["Alice", "Bob"],
year=2023,
abstract=abstract,
)
def _fake_embedder(dim: int = 1024) -> MagicMock:
"""Return a mock Embedder whose encode_dense returns deterministic output."""
embedder = MagicMock()
embedder.dim = dim
def encode_dense(texts: list[str]) -> np.ndarray:
return np.ones((len(texts), dim), dtype=np.float32)
embedder.encode_dense.side_effect = encode_dense
return embedder
# ---------------------------------------------------------------------------
# 1. test_ingest_paper_basic
# ---------------------------------------------------------------------------
def test_ingest_paper_basic() -> None:
"""fetch_paper returns a Paper; DB upsert is called; IngestResult has correct counts."""
paper = _make_paper()
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
# OpenAlex returns citing_id=openalex_id; ingest.py must rewrite to paper.id (D-05c).
raw_api_citations = [
Citation(citing_id=paper.openalex_id or "", cited_id="W999"),
]
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper) as mock_fetch,
patch("codex.ingest.openalex.fetch_citations", return_value=raw_api_citations),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
result = ingest_paper(paper.id)
mock_fetch.assert_called_once_with(paper.id)
# Paper upsert must have been issued
assert mock_conn.execute.call_count >= 1
first_sql: str = mock_conn.execute.call_args_list[0][0][0]
assert "INSERT INTO papers" in first_sql
assert "ON CONFLICT" in first_sql
assert isinstance(result, IngestResult)
assert result.paper_id == paper.id
assert result.chunks_upserted == 0
assert result.citations_upserted == len(raw_api_citations)
# D-05c regression: citing_id must be rewritten to paper.id (DOI), not openalex_id.
mock_cursor = mock_conn.cursor.return_value.__enter__.return_value
cit_call = mock_cursor.executemany.call_args_list[0]
inserted_rows: list[tuple[str, str, Any]] = cit_call[0][1]
assert inserted_rows[0][0] == paper.id, (
f"citing_id must be paper.id='{paper.id}', not openalex_id='{paper.openalex_id}'"
)
def test_ingest_canonicalises_paper_id_to_caller_arg() -> None:
"""papers.id is the caller-supplied id, not OpenAlex's URL-form id (audit C-7).
Real OpenAlex returns the doi/id as a full URL; ingest must key the paper on
the bare id the caller passed (matching ingest scripts and CLI lookups), and
keep the OpenAlex work id only as openalex_id.
"""
oa_paper = Paper(
id="https://doi.org/10.48550/arxiv.math/0603097", # URL form, as OpenAlex emits
title="A Test Paper",
openalex_id="https://openalex.org/W4301005794",
authors=["Alice", "Bob"],
year=2008,
abstract="Test abstract.",
)
mock_conn = MagicMock()
with (
patch("codex.ingest.openalex.fetch_paper", return_value=oa_paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
result = ingest_paper("math/0603097")
papers_params = mock_conn.execute.call_args_list[0][0][1]
assert papers_params["id"] == "math/0603097", "papers.id must be the caller's bare id"
assert papers_params["openalex_id"] == "https://openalex.org/W4301005794"
assert result.paper_id == "math/0603097"
# ---------------------------------------------------------------------------
# 2. test_ingest_paper_source_tex
# ---------------------------------------------------------------------------
def test_ingest_paper_source_tex(tmp_path: Any) -> None:
"""``.tex`` source → latex_to_text, chunk_text, embed, chunks inserted in DB."""
tex_file = tmp_path / "paper.tex"
tex_file.write_text(r"\section{Introduction} Hello world. This is a test paper.")
paper = _make_paper()
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
fake_chunks = ["chunk one text here.", "chunk two text here."]
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
patch("codex.parsing.tex.latex_to_text", return_value="Hello world. Intro text."),
patch("codex.parsing.tex.chunk_text", return_value=fake_chunks),
patch("codex.ingest.filter_chunks", side_effect=lambda chunks, settings: chunks),
):
result = ingest_paper(paper.id, source_path=str(tex_file))
# DELETE + INSERT for chunks
sql_calls = [str(c[0][0]) for c in mock_conn.execute.call_args_list]
assert any("DELETE FROM chunks" in sql for sql in sql_calls)
# executemany is called on the cursor, not on the connection directly
mock_cursor = mock_conn.cursor.return_value.__enter__.return_value
assert mock_cursor.executemany.call_count >= 1
chunk_insert_call = mock_cursor.executemany.call_args_list[0]
chunk_sql: str = chunk_insert_call[0][0]
assert "INSERT INTO chunks" in chunk_sql
assert result.chunks_upserted == len(fake_chunks)
# ---------------------------------------------------------------------------
# 3. test_ingest_paper_source_pdf
# ---------------------------------------------------------------------------
def test_ingest_paper_source_pdf(tmp_path: Any) -> None:
"""``.pdf`` source → pdf_to_markdown, chunk_text, embed, grobid refs, chunks + citations."""
pdf_file = tmp_path / "paper.pdf"
pdf_file.write_bytes(b"%PDF-1.4 fake content")
paper = _make_paper()
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
fake_chunks = ["pdf chunk one.", "pdf chunk two."]
grobid_refs = [
{"title": "Ref A", "authors": "X", "year": "2020", "doi": "10.1/ref_a", "arxiv_id": ""},
{"title": "Ref B", "authors": "Y", "year": "2021", "doi": "", "arxiv_id": "2101.00001"},
]
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
patch("codex.parsing.nougat.pdf_to_markdown", return_value="pdf markdown text."),
patch("codex.parsing.tex.chunk_text", return_value=fake_chunks),
patch("codex.parsing.grobid.extract_references", return_value=grobid_refs),
patch("codex.ingest.filter_chunks", side_effect=lambda chunks, settings: chunks),
):
result = ingest_paper(paper.id, source_path=str(pdf_file))
# Chunks were inserted
assert result.chunks_upserted == len(fake_chunks)
# GROBID refs were merged into citations
# Two grobid refs → both have either doi or arxiv_id → 2 citations
assert result.citations_upserted == 2
# executemany is called on the cursor for chunks + citations (at least 2 calls total)
mock_cursor = mock_conn.cursor.return_value.__enter__.return_value
assert mock_cursor.executemany.call_count >= 2
# ---------------------------------------------------------------------------
# 4. test_ingest_paper_not_found
# ---------------------------------------------------------------------------
def test_ingest_paper_not_found() -> None:
"""OpenAlex returns None for unknown ID → ValueError raised."""
with (
patch("codex.ingest.openalex.fetch_paper", return_value=None),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
pytest.raises(ValueError, match="Paper not found"),
):
ingest_paper("10.9999/unknown-doi")
# ---------------------------------------------------------------------------
# 5. test_ingest_paper_idempotent
# ---------------------------------------------------------------------------
def test_ingest_paper_idempotent() -> None:
"""Second call with the same paper_id overwrites without error (no unique violations)."""
paper = _make_paper()
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
result1 = ingest_paper(paper.id)
result2 = ingest_paper(paper.id)
assert result1.paper_id == result2.paper_id == paper.id
# ON CONFLICT … DO UPDATE means no error on second call
all_execute_calls = mock_conn.execute.call_args_list
paper_upserts = [c for c in all_execute_calls if "INSERT INTO papers" in str(c[0][0])]
# Two calls → two paper upserts
assert len(paper_upserts) == 2
# ---------------------------------------------------------------------------
# 6. test_ingest_paper_arxiv_s2_fallback
# ---------------------------------------------------------------------------
def test_ingest_paper_arxiv_s2_fallback() -> None:
"""If OpenAlex returns None for an arXiv ID, S2 fallback produces a stub Paper."""
paper_id = "2301.07041"
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
with (
patch("codex.ingest.openalex.fetch_paper", return_value=None),
patch("codex.ingest.semanticscholar.fetch_references", return_value=[]) as mock_s2,
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
result = ingest_paper(paper_id)
# S2 was consulted as fallback
mock_s2.assert_called()
assert result.paper_id == paper_id
# ---------------------------------------------------------------------------
# 7. test_ingest_paper_no_abstract_uses_zero_vector
# ---------------------------------------------------------------------------
def test_ingest_paper_no_abstract_uses_zero_vector() -> None:
"""A paper without an abstract gets a zero embedding vector (no model call)."""
paper = _make_paper(abstract=None)
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
fake_emb = _fake_embedder(dim=1024)
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=fake_emb),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
ingest_paper(paper.id)
# encode_dense must NOT have been called (no abstract → zero vector shortcut)
fake_emb.encode_dense.assert_not_called()
# Verify zero vector was passed in the paper upsert
upsert_call = mock_conn.execute.call_args_list[0]
params: dict[str, Any] = upsert_call[0][1]
emb: list[float] = params["abstract_emb"]
assert len(emb) == 1024
assert all(v == 0.0 for v in emb)
# ---------------------------------------------------------------------------
# 8. F-09 Rich Parsing tests
# ---------------------------------------------------------------------------
def test_ingest_paper_rich_calls_formula_and_figure_parsers(tmp_path: Any) -> None:
"""When rich=True and source is PDF, formula + figure parsers are invoked."""
import codex.parsing.figures as _figures_mod
import codex.parsing.mathpix as _mathpix_mod
pdf_file = tmp_path / "paper.pdf"
pdf_file.write_bytes(b"%PDF-1.4 fake content")
paper = _make_paper()
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
# Extractors derive paper_id from the PDF filename stem ("paper"), which differs
# from papers.id — ingest must insert papers.id, not this stem (audit C-10).
fake_formula = FormulaChunk(paper_id="paper", page=1, raw_latex=r"\alpha", context="test")
fake_figure = FigureChunk(
paper_id="paper", page=1, image_path="/tmp/fig.png", caption="Figure 1."
)
mock_settings = MagicMock()
mock_settings.figures_dir = str(tmp_path / "figures")
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
patch("codex.parsing.nougat.pdf_to_markdown", return_value=""),
patch("codex.parsing.tex.chunk_text", return_value=[]),
patch("codex.parsing.grobid.extract_references", return_value=[]),
patch.object(
_mathpix_mod, "extract_formulas", return_value=[fake_formula]
) as mock_formulas,
patch.object(_figures_mod, "extract_figures", return_value=[fake_figure]) as mock_figures,
patch("codex.ingest.get_settings", return_value=mock_settings),
):
result = ingest_paper(paper.id, source_path=str(pdf_file), rich=True)
mock_formulas.assert_called_once()
mock_figures.assert_called_once()
assert result.formulas_upserted == 1
assert result.figures_upserted == 1
# C-10: formula/figure rows must be keyed on papers.id, not the extractor's
# filename-stem paper_id, or the FK to papers(id) is violated.
cursor = mock_conn.cursor.return_value.__enter__.return_value
inserted_paper_ids = {
row[0] for call in cursor.executemany.call_args_list for row in call[0][1]
}
assert inserted_paper_ids == {paper.id}, (
f"formula/figure inserts must use papers.id, got {inserted_paper_ids}"
)
def test_ingest_paper_no_rich_skips_parsers(tmp_path: Any) -> None:
"""When rich=False (default), formula and figure parsers are NOT called."""
import codex.parsing.figures as _figures_mod
import codex.parsing.mathpix as _mathpix_mod
pdf_file = tmp_path / "paper.pdf"
pdf_file.write_bytes(b"%PDF-1.4 fake content")
paper = _make_paper()
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
patch("codex.parsing.nougat.pdf_to_markdown", return_value=""),
patch("codex.parsing.tex.chunk_text", return_value=[]),
patch("codex.parsing.grobid.extract_references", return_value=[]),
patch.object(_mathpix_mod, "extract_formulas", return_value=[]) as mock_formulas,
patch.object(_figures_mod, "extract_figures", return_value=[]) as mock_figures,
):
result = ingest_paper(paper.id, source_path=str(pdf_file), rich=False)
mock_formulas.assert_not_called()
mock_figures.assert_not_called()
assert result.formulas_upserted == 0
assert result.figures_upserted == 0
def test_ingest_paper_rich_requires_pdf_source(tmp_path: Any) -> None:
"""When rich=True but source is a .tex file, formula/figure parsers are NOT called."""
import codex.parsing.figures as _figures_mod
import codex.parsing.mathpix as _mathpix_mod
tex_file = tmp_path / "paper.tex"
tex_file.write_text(r"\section{Intro}")
paper = _make_paper()
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
patch("codex.parsing.tex.latex_to_text", return_value=""),
patch("codex.parsing.tex.chunk_text", return_value=[]),
patch.object(_mathpix_mod, "extract_formulas", return_value=[]) as mock_formulas,
patch.object(_figures_mod, "extract_figures", return_value=[]) as mock_figures,
):
result = ingest_paper(paper.id, source_path=str(tex_file), rich=True)
mock_formulas.assert_not_called()
mock_figures.assert_not_called()
assert result.formulas_upserted == 0
assert result.figures_upserted == 0
# ---------------------------------------------------------------------------
# D-04 — bibkey auto-generation heuristic
# ---------------------------------------------------------------------------
def test_make_bibkey_single_author() -> None:
paper = Paper(id="doi:10.1/x", title="T", authors=["John Smith"], year=2023)
assert _make_bibkey(paper) == "Smith2023"
def test_make_bibkey_two_authors() -> None:
paper = Paper(id="doi:10.1/x", title="T", authors=["Alice Foo", "Bob Bar"], year=2015)
assert _make_bibkey(paper) == "FooBar2015"
def test_make_bibkey_three_authors() -> None:
paper = Paper(
id="doi:10.1/x",
title="T",
authors=["Alexander Bobenko", "Ulrich Pinkall", "Boris Springborn"],
year=2015,
)
assert _make_bibkey(paper) == "BobenkoPinkallSpringborn2015"
def test_make_bibkey_many_authors_uses_etal() -> None:
paper = Paper(
id="doi:10.1/x",
title="T",
authors=["A One", "B Two", "C Three", "D Four"],
year=2020,
)
assert _make_bibkey(paper) == "OneEtAl2020"
def test_make_bibkey_no_authors_returns_none() -> None:
paper = Paper(id="doi:10.1/x", title="T", authors=[], year=2023)
assert _make_bibkey(paper) is None
def test_make_bibkey_no_year_returns_none() -> None:
paper = Paper(id="doi:10.1/x", title="T", authors=["John Smith"], year=None)
assert _make_bibkey(paper) is None
def test_ingest_sets_bibkey_when_none(tmp_path: Any) -> None:
"""ingest_paper auto-generates bibkey when paper has none (D-04)."""
paper = _make_paper() # bibkey=None (default from _make_paper)
mock_conn = MagicMock()
mock_conn.execute = MagicMock()
mock_conn.executemany = MagicMock()
mock_conn.commit = MagicMock()
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
ingest_paper(paper.id)
upsert_params: dict[str, Any] = mock_conn.execute.call_args_list[0][0][1]
assert upsert_params["bibkey"] is not None, "bibkey must be auto-generated when paper has none"
# authors=["Alice", "Bob"], year=2023 → 2 authors ≤ 3 → concat last tokens
assert upsert_params["bibkey"] == "AliceBob2023"
def test_ingest_disambiguates_bibkey_on_unique_violation() -> None:
"""A papers.bibkey UNIQUE collision is retried with an a/b/c suffix (audit C-15)."""
paper = _make_paper() # bibkey auto-generates to "AliceBob2023"
mock_conn = MagicMock()
state = {"first": True}
def execute_side_effect(sql: str, params: Any = None) -> MagicMock:
# The first papers INSERT collides on the bibkey; the retry must succeed.
if "INSERT INTO papers" in sql and state["first"]:
state["first"] = False
raise psycopg.errors.UniqueViolation(
'duplicate key value violates unique constraint "papers_bibkey_key"'
)
return MagicMock()
mock_conn.execute.side_effect = execute_side_effect
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
ingest_paper(paper.id)
papers_inserts = [
c for c in mock_conn.execute.call_args_list if "INSERT INTO papers" in c[0][0]
]
assert len(papers_inserts) == 2, "expected one retry after the bibkey collision"
assert papers_inserts[0][0][1]["bibkey"] == "AliceBob2023"
assert papers_inserts[1][0][1]["bibkey"] == "AliceBob2023a"
mock_conn.rollback.assert_called_once()
def test_ingest_openalex_collision_raises_clear_error() -> None:
"""An openalex_id UNIQUE collision raises a clear, actionable error (audit C-15).
Unlike a bibkey clash (auto-suffixed), a shared openalex_id means two papers
claim the same OpenAlex work — a real data conflict that must surface, not the
raw psycopg error.
"""
paper = _make_paper()
mock_conn = MagicMock()
def execute_side_effect(sql: str, params: Any = None) -> MagicMock:
if "INSERT INTO papers" in sql:
raise psycopg.errors.UniqueViolation(
'duplicate key value violates unique constraint "papers_openalex_id_key"'
)
return MagicMock()
mock_conn.execute.side_effect = execute_side_effect
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
pytest.raises(ValueError, match="openalex_id"),
):
ingest_paper(paper.id)
# ---------------------------------------------------------------------------
def test_ingest_result_has_formula_figure_counts() -> None:
"""IngestResult has formulas_upserted and figures_upserted fields with default 0."""
result = IngestResult(paper_id="test", chunks_upserted=0, citations_upserted=0)
assert result.formulas_upserted == 0
assert result.figures_upserted == 0
result2 = IngestResult(
paper_id="test",
chunks_upserted=5,
citations_upserted=3,
formulas_upserted=10,
figures_upserted=2,
)
assert result2.formulas_upserted == 10
assert result2.figures_upserted == 2
Reference in New Issue View Git Blame Copy Permalink