#!/usr/bin/env python3 """ Tests for scripts/graph_query.py — Graph Query Engine. """ import json import sys import tempfile from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent)) from graph_query import load_index, build_adjacency, query_neighbors, query_path, query_subgraph, query_stats def make_index(facts: list[dict], tmp_dir: Path) -> Path: index = { "version": 1, "last_updated": "2026-04-13T20:00:00Z", "total_facts": len(facts), "facts": facts, } path = tmp_dir / "index.json" with open(path, "w") as f: json.dump(index, f) return path def test_neighbors(): """Neighbor query returns directly connected facts.""" facts = [ {"id": "a", "fact": "A", "category": "fact", "related": ["b", "c"]}, {"id": "b", "fact": "B", "category": "fact", "related": ["a"]}, {"id": "c", "fact": "C", "category": "fact", "related": ["a"]}, {"id": "d", "fact": "D", "category": "fact", "related": []}, ] adj, id_to_fact = build_adjacency(facts) result = query_neighbors("a", adj, id_to_fact) neighbor_ids = {n["id"] for n in result["neighbors"]} assert neighbor_ids == {"b", "c"}, f"Expected b,c got {neighbor_ids}" assert result["count"] == 2 print("PASS: neighbors") def test_path_found(): """Path query finds shortest path.""" facts = [ {"id": "a", "fact": "A", "related": ["b"]}, {"id": "b", "fact": "B", "related": ["a", "c"]}, {"id": "c", "fact": "C", "related": ["b", "d"]}, {"id": "d", "fact": "D", "related": ["c"]}, ] adj, id_to_fact = build_adjacency(facts) result = query_path("a", "d", adj) assert result["path"] == ["a", "b", "c", "d"], f"Got path {result['path']}" assert result["length"] == 3 print("PASS: path_found") def test_path_not_found(): """Path query returns error when no path exists.""" facts = [ {"id": "a", "fact": "A", "related": ["b"]}, {"id": "b", "fact": "B", "related": ["a"]}, {"id": "c", "fact": "C", "related": ["d"]}, {"id": "d", "fact": "D", "related": ["c"]}, ] adj, id_to_fact = build_adjacency(facts) result = query_path("a", "c", adj, max_hops=5) assert result["path"] is None assert "error" in result print("PASS: path_not_found") def test_subgraph_extraction(): """Subgraph extraction returns nodes within depth.""" facts = [ {"id": "a", "fact": "A", "related": ["b", "c"]}, {"id": "b", "fact": "B", "related": ["a", "d"]}, {"id": "c", "fact": "C", "related": ["a"]}, {"id": "d", "fact": "D", "related": ["b", "e"]}, {"id": "e", "fact": "E", "related": ["d"]}, ] adj, id_to_fact = build_adjacency(facts) result = query_subgraph("a", adj, id_to_fact, depth=1) node_ids = {n["id"] for n in result["nodes"]} assert node_ids == {"a", "b", "c"}, f"Got {node_ids}" assert result["node_count"] == 3 print("PASS: subgraph_depth1") def test_subgraph_depth2(): """Depth-2 subgraph includes further nodes.""" facts = [ {"id": "a", "fact": "A", "related": ["b"]}, {"id": "b", "fact": "B", "related": ["a", "c"]}, {"id": "c", "fact": "C", "related": ["b", "d"]}, {"id": "d", "fact": "D", "related": ["c"]}, ] adj, id_to_fact = build_adjacency(facts) result = query_subgraph("a", adj, id_to_fact, depth=2) node_ids = {n["id"] for n in result["nodes"]} assert node_ids == {"a", "b", "c"}, f"Got {node_ids}" print("PASS: subgraph_depth2") def test_stats(): """Statistics query returns graph metrics.""" facts = [ {"id": "a", "fact": "A", "related": ["b"]}, {"id": "b", "fact": "B", "related": ["a", "c"]}, {"id": "c", "fact": "C", "related": ["b"]}, ] adj, id_to_fact = build_adjacency(facts) result = query_stats(adj, id_to_fact) assert result["statistics"]["total_facts"] == 3 assert result["statistics"]["total_edges"] == 2 # undirected double-counted /2 assert result["statistics"]["average_degree"] > 0 print("PASS: stats") def test_cli_integration(): """CLI produces valid JSON with correct query types.""" with tempfile.TemporaryDirectory() as tmp: import subprocess as sp tmp_dir = Path(tmp) facts = [ {"id": "x", "fact": "X", "related": ["y"]}, {"id": "y", "fact": "Y", "related": ["x", "z"]}, {"id": "z", "fact": "Z", "related": ["y"]}, ] index_path = make_index(facts, tmp_dir) knowledge_dir = index_path.parent script_path = Path(__file__).resolve().parent / "graph_query.py" result = sp.run( [sys.executable, str(script_path), "neighbors", "x", "--knowledge-dir", str(knowledge_dir)], capture_output=True, text=True, cwd=str(tmp_dir) ) assert result.returncode == 0, f"neighbors failed: {result.stderr}" out = json.loads(result.stdout) assert out["query"] == "neighbors" assert out["fact_id"] == "x" assert out["count"] == 1 result = sp.run( [sys.executable, str(script_path), "path", "x", "z", "--knowledge-dir", str(knowledge_dir)], capture_output=True, text=True, cwd=str(tmp_dir) ) assert result.returncode == 0, f"path failed: {result.stderr}" out = json.loads(result.stdout) assert out["path"] == ["x", "y", "z"] print("PASS: cli_integration") if __name__ == "__main__": test_neighbors() test_path_found() test_path_not_found() test_subgraph_extraction() test_subgraph_depth2() test_stats() test_cli_integration() print("\nAll graph_query tests passed!")