codex-py/tests/provenance/test_provenance.py

"""Tests for codex.provenance — @cite scan, code_links CRUD, bib export.

All DB access is mocked; no real PostgreSQL server required.
"""

from __future__ import annotations

from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch

from codex.models import CodeLink

# ---------------------------------------------------------------------------
# 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 _dict_row(**kwargs: Any) -> dict[str, Any]:
    return dict(kwargs)


# ---------------------------------------------------------------------------
# 1. test_scan_cite_tags
# ---------------------------------------------------------------------------


def test_scan_cite_tags(tmp_path: Path) -> None:
    """A .cpp file with @cite produces correct dict entries."""
    from codex.provenance import scan_cite_tags

    src_file = tmp_path / "algo.cpp"
    src_file.write_text(
        "// Helper function\n"
        "void compute_hodge() {\n"
        "    // @cite somepaper2023\n"
        "    do_stuff();\n"
        "}\n",
        encoding="utf-8",
    )

    results = scan_cite_tags(str(tmp_path))

    assert len(results) == 1
    hit = results[0]
    assert hit["bibkey"] == "somepaper2023"
    assert hit["line"] == "3"
    assert hit["file"] == "algo.cpp"
    # symbol extracted from "compute_hodge()" on the line above
    assert hit["symbol"] == "compute_hodge"


# ---------------------------------------------------------------------------
# 2. test_scan_cite_tags_no_files
# ---------------------------------------------------------------------------


def test_scan_cite_tags_no_files(tmp_path: Path) -> None:
    """Empty directory returns an empty list."""
    from codex.provenance import scan_cite_tags

    result = scan_cite_tags(str(tmp_path))
    assert result == []


# ---------------------------------------------------------------------------
# 3. test_add_code_link
# ---------------------------------------------------------------------------


def test_add_code_link() -> None:
    """Mock conn, verify INSERT is issued and returned CodeLink has id from RETURNING."""
    from codex.provenance import add_code_link

    now = datetime(2026, 6, 5, tzinfo=UTC)
    mock_conn = MagicMock()
    mock_conn.execute.return_value.fetchone.return_value = _dict_row(id=42, added_at=now)

    with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)):
        link = add_code_link(
            symbol="dec::hodge_star",
            paper_id="2301.07041",
            role="implements Thm 3.2",
            note="core algorithm",
        )

    assert isinstance(link, CodeLink)
    assert link.id == 42
    assert link.added_at == now
    assert link.symbol == "dec::hodge_star"
    assert link.paper_id == "2301.07041"
    assert link.role == "implements Thm 3.2"
    assert link.note == "core algorithm"

    call_args = mock_conn.execute.call_args
    sql: str = call_args[0][0]
    assert "INSERT INTO code_links" in sql
    assert "RETURNING" in sql

    mock_conn.commit.assert_called_once()


# ---------------------------------------------------------------------------
# 4. test_list_code_links_all
# ---------------------------------------------------------------------------


def test_list_code_links_all() -> None:
    """No filter → all rows returned as CodeLink list."""
    from codex.provenance import list_code_links

    now = datetime(2026, 6, 5, tzinfo=UTC)
    mock_conn = MagicMock()
    mock_conn.execute.return_value.fetchall.return_value = [
        _dict_row(id=1, symbol="foo::bar", paper_id="W111", role="uses", note=None, added_at=now),
        _dict_row(id=2, symbol="baz::qux", paper_id=None, role=None, note="ref", added_at=now),
    ]

    with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)):
        links = list_code_links()

    assert len(links) == 2
    assert all(isinstance(lnk, CodeLink) for lnk in links)
    assert links[0].symbol == "foo::bar"
    assert links[0].id == 1
    assert links[1].symbol == "baz::qux"
    assert links[1].paper_id is None

    # No WHERE clause in the SQL
    call_args = mock_conn.execute.call_args
    sql: str = call_args[0][0]
    assert "WHERE" not in sql


# ---------------------------------------------------------------------------
# 5. test_list_code_links_filtered
# ---------------------------------------------------------------------------


def test_list_code_links_filtered() -> None:
    """Filter by paper_id → WHERE clause applied and only matching rows returned."""
    from codex.provenance import list_code_links

    now = datetime(2026, 6, 5, tzinfo=UTC)
    mock_conn = MagicMock()
    mock_conn.execute.return_value.fetchall.return_value = [
        _dict_row(
            id=3,
            symbol="some::func",
            paper_id="2301.07041",
            role=None,
            note=None,
            added_at=now,
        ),
    ]

    with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)):
        links = list_code_links(paper_id="2301.07041")

    assert len(links) == 1
    assert links[0].paper_id == "2301.07041"

    call_args = mock_conn.execute.call_args
    sql: str = call_args[0][0]
    params: dict[str, Any] = call_args[0][1]
    assert "WHERE paper_id" in sql
    assert params["paper_id"] == "2301.07041"


# ---------------------------------------------------------------------------
# 6. test_export_bib_format
# ---------------------------------------------------------------------------


def test_export_bib_format() -> None:
    """Mock returns one paper row; verify valid BibTeX string output."""
    from codex.provenance import export_bib

    mock_conn = MagicMock()
    mock_conn.execute.return_value.fetchall.return_value = [
        _dict_row(
            id="2301.07041",
            bibkey="smith2023conformal",
            title="Conformal Approaches to Laplacian Eigenmaps",
            authors=["Alice Smith", "Bob Jones"],
            year=2023,
        ),
    ]

    with patch("codex.provenance.get_conn", side_effect=_make_conn_cm(mock_conn)):
        bib = export_bib()

    assert "@article{smith2023conformal," in bib
    assert "title  = {Conformal Approaches to Laplacian Eigenmaps}" in bib
    assert "author = {Alice Smith and Bob Jones}" in bib
    assert "year   = {2023}" in bib
    assert "note   = {2301.07041}" in bib