kg-scr/tests/test_graph.py

from __future__ import annotations

from pathlib import Path

import pytest

nx = pytest.importorskip("networkx", reason="kg extra not installed")

from kg_ocr.export import export_graphml, export_json, load_json  # noqa: E402
from kg_ocr.graph import (  # noqa: E402
    build_from_directory,
    build_graph,
    detect_anomalies,
    find_chunk_files,
    parse_chunk_file,
    summary,
    top_citations,
    top_cooccurrences,
    top_entities,
)

CHUNK_MD = """---
source_path: /shots/SCR-001.png
source_hash: aaaabbbbccccdddd
timestamp: '2026-07-19T10:00:00+00:00'
ocr_engine: paddleocr
ocr_confidence_mean: 0.92
language: en
detected_entities:
- BRCA1
- PARP
citations_found:
- doi:10.1038/nature12345
chunk_index: 0
total_chunks: 1
---

# Screenshot: SCR-001.png

## Extracted Text

BRCA1 and PARP inhibition in ovarian cancer.

## Citations Found

| Type | Identifier |
"""

CHUNK_MD_2 = """---
source_path: /shots/SCR-002.png
source_hash: eeeeffff00001111
timestamp: '2026-07-19T10:01:00+00:00'
ocr_engine: tesseract
ocr_confidence_mean: 0.2
language: en
detected_entities:
- BRCA1
citations_found: []
chunk_index: 0
total_chunks: 1
---

# Screenshot: SCR-002.png

## Extracted Text

BRCA1 western blot.
"""


@pytest.fixture()
def chunk_dir(tmp_path: Path) -> Path:
    run_dir = tmp_path / "out" / "2026-07-19" / "run_001"
    run_dir.mkdir(parents=True)
    (run_dir / "SCR-001_aaaabbbbcccc.md").write_text(CHUNK_MD, encoding="utf-8")
    (run_dir / "SCR-002_eeeeffff0000.md").write_text(CHUNK_MD_2, encoding="utf-8")
    return tmp_path / "out"


def test_parse_chunk_file(chunk_dir: Path) -> None:
    path = next(chunk_dir.rglob("SCR-001_*.md"))
    record = parse_chunk_file(path)
    assert record is not None
    assert record.source_hash == "aaaabbbbccccdddd"
    assert record.ocr_confidence_mean == pytest.approx(0.92)
    assert record.entities == ["BRCA1", "PARP"]
    assert record.citations == ["doi:10.1038/nature12345"]
    assert "BRCA1 and PARP" in record.text


def test_parse_chunk_file_rejects_plain_markdown(tmp_path: Path) -> None:
    plain = tmp_path / "notes.md"
    plain.write_text("# just notes\n", encoding="utf-8")
    assert parse_chunk_file(plain) is None


def test_find_chunk_files_skips_consolidated(chunk_dir: Path) -> None:
    (chunk_dir / "all_ocr.md").write_text("# consolidated\n", encoding="utf-8")
    names = [p.name for p in find_chunk_files(chunk_dir)]
    assert "all_ocr.md" not in names
    assert len(names) == 2


def test_build_graph_structure(chunk_dir: Path) -> None:
    graph = build_from_directory(chunk_dir)
    counts = summary(graph)

    assert counts["nodes:document"] == 2
    assert counts["nodes:chunk"] == 2
    assert counts["nodes:entity"] == 2  # brca1, parp (normalized)
    assert counts["nodes:citation"] == 1
    assert counts["edges:CONTAINS"] == 2
    assert counts["edges:MENTIONS"] == 3
    assert counts["edges:CITES"] == 1
    assert counts["edges:CO_OCCURS"] == 1  # brca1 <-> parp

    assert graph.has_edge("doc:aaaabbbbccccdddd", "chunk:aaaabbbbccccdddd:0")
    assert graph.has_edge("chunk:aaaabbbbccccdddd:0", "entity:brca1")


def test_top_entities_and_cooccurrence(chunk_dir: Path) -> None:
    graph = build_from_directory(chunk_dir)
    top = top_entities(graph)
    assert top[0] == ("BRCA1", 2)
    pairs = top_cooccurrences(graph)
    assert pairs == [("BRCA1", "PARP", 1)]


def test_top_citations(chunk_dir: Path) -> None:
    graph = build_from_directory(chunk_dir)
    rows = top_citations(graph)
    assert rows == [("doi", "10.1038/nature12345", 1)]


def test_detect_anomalies(chunk_dir: Path) -> None:
    graph = build_from_directory(chunk_dir)
    kinds = {a.kind for a in detect_anomalies(graph)}
    # SCR-002 has confidence 0.2 (< 0.4) -> low_confidence
    assert "low_confidence" in kinds


def test_json_round_trip(chunk_dir: Path, tmp_path: Path) -> None:
    graph = build_from_directory(chunk_dir)
    out = export_json(graph, tmp_path / "g.json")
    restored = load_json(out)
    assert summary(restored) == summary(graph)
    assert restored.nodes["entity:brca1"]["mentions"] == 2


def test_graphml_export(chunk_dir: Path, tmp_path: Path) -> None:
    graph = build_from_directory(chunk_dir)
    out = export_graphml(graph, tmp_path / "g.graphml")
    content = out.read_text(encoding="utf-8")
    assert "graphml" in content


def test_empty_directory(tmp_path: Path) -> None:
    graph = build_from_directory(tmp_path)
    assert summary(graph) == {}


def test_build_graph_from_records_directly() -> None:
    from kg_ocr.graph import ChunkRecord

    record = ChunkRecord(
        source_hash="x" * 16,
        source_path="/a.png",
        chunk_index=0,
        total_chunks=1,
        text="hello",
        entities=["TP53", "tp53"],  # dedupe via normalization
        citations=["pmid:PMID:12345678"],
    )
    graph = build_graph([record])
    counts = summary(graph)
    assert counts["nodes:entity"] == 1
    assert graph.nodes["entity:tp53"]["mentions"] == 2
    assert graph.nodes["citation:PMID:12345678"]["citation_type"] == "pmid"


def test_kg_cli_build_and_stats(chunk_dir: Path, tmp_path: Path) -> None:
    from typer.testing import CliRunner

    from ocr_pipeline.cli import app

    runner = CliRunner()
    save = tmp_path / "g.json"

    result = runner.invoke(app, ["kg", "build", "-d", str(chunk_dir), "-s", str(save)])
    assert result.exit_code == 0, result.output
    assert save.is_file()

    result = runner.invoke(app, ["kg", "stats", str(save)])
    assert result.exit_code == 0, result.output
    assert "BRCA1" in result.output
    assert "low_confidence" in result.output

    out_graphml = tmp_path / "g.graphml"
    result = runner.invoke(
        app, ["kg", "export", str(save), "-f", "graphml", "-o", str(out_graphml)]
    )
    assert result.exit_code == 0, result.output
    assert out_graphml.is_file()


def test_chunks_for_entity(chunk_dir: Path) -> None:
    from kg_ocr.graph import chunks_for_entity

    graph = build_from_directory(chunk_dir)
    chunks = chunks_for_entity(graph, "brca1")  # case-insensitive
    assert len(chunks) == 2
    assert {c["source_path"] for c in chunks} == {"/shots/SCR-001.png", "/shots/SCR-002.png"}
    assert chunks_for_entity(graph, "NONEXISTENT") == []


def test_chunks_for_citation(chunk_dir: Path) -> None:
    from kg_ocr.graph import chunks_for_citation

    graph = build_from_directory(chunk_dir)
    chunks = chunks_for_citation(graph, "10.1038/nature12345")
    assert len(chunks) == 1
    assert "BRCA1 and PARP" in chunks[0]["text"]
    assert chunks_for_citation(graph, "nope") == []


def test_related_entities(chunk_dir: Path) -> None:
    from kg_ocr.graph import related_entities

    graph = build_from_directory(chunk_dir)
    assert related_entities(graph, "BRCA1") == [("PARP", 1)]
    assert related_entities(graph, "PARP") == [("BRCA1", 1)]
    assert related_entities(graph, "NONEXISTENT") == []


def test_expand_context(chunk_dir: Path) -> None:
    from kg_ocr.graph import expand_context

    graph = build_from_directory(chunk_dir)
    chunks = expand_context(graph, "PARP")
    # PARP's own chunk + BRCA1 neighbor chunks (SCR-001 deduped)
    assert len(chunks) == 2
    via = {c["via_entity"] for c in chunks}
    assert "PARP" in via
    assert "BRCA1" in via


def test_kg_cli_query(chunk_dir: Path, tmp_path: Path) -> None:
    from typer.testing import CliRunner

    from ocr_pipeline.cli import app

    runner = CliRunner()
    save = tmp_path / "g.json"
    result = runner.invoke(app, ["kg", "build", "-d", str(chunk_dir), "-s", str(save)])
    assert result.exit_code == 0, result.output

    result = runner.invoke(app, ["kg", "query", str(save), "--entity", "PARP"])
    assert result.exit_code == 0, result.output
    assert "SCR-001.png" in result.output
    assert "Related entities" in result.output

    result = runner.invoke(app, ["kg", "query", str(save), "--citation", "10.1038/nature12345"])
    assert result.exit_code == 0, result.output
    assert "SCR-001.png" in result.output

    result = runner.invoke(app, ["kg", "query", str(save)])
    assert result.exit_code == 2  # neither --entity nor --citation