"""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 pytest from codex.ingest import IngestResult, ingest_paper from codex.models import Citation, 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() 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=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(api_citations) # --------------------------------------------------------------------------- # 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), ): 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), ): 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)