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