Claude-Code-Workflow/codex-lens/tests/test_graph_expansion.py

import sqlite3
import tempfile
from pathlib import Path

import pytest

from codexlens.config import Config
from codexlens.entities import CodeRelationship, RelationshipType, SearchResult, Symbol
from codexlens.search.chain_search import ChainSearchEngine, SearchOptions
from codexlens.search.graph_expander import GraphExpander
from codexlens.storage.dir_index import DirIndexStore
from codexlens.storage.index_tree import _compute_graph_neighbors
from codexlens.storage.path_mapper import PathMapper
from codexlens.storage.registry import RegistryStore


@pytest.fixture()
def temp_paths() -> Path:
    tmpdir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
    root = Path(tmpdir.name)
    yield root
    try:
        tmpdir.cleanup()
    except (PermissionError, OSError):
        pass


def _create_index_with_neighbors(root: Path) -> tuple[PathMapper, Path, Path]:
    project_root = root / "project"
    project_root.mkdir(parents=True, exist_ok=True)

    index_root = root / "indexes"
    mapper = PathMapper(index_root=index_root)
    index_db_path = mapper.source_to_index_db(project_root)
    index_db_path.parent.mkdir(parents=True, exist_ok=True)

    content = "\n".join(
        [
            "def a():",
            "    b()",
            "",
            "def b():",
            "    c()",
            "",
            "def c():",
            "    return 1",
            "",
        ]
    )
    file_path = project_root / "module.py"
    file_path.write_text(content, encoding="utf-8")

    symbols = [
        Symbol(name="a", kind="function", range=(1, 2), file=str(file_path)),
        Symbol(name="b", kind="function", range=(4, 5), file=str(file_path)),
        Symbol(name="c", kind="function", range=(7, 8), file=str(file_path)),
    ]
    relationships = [
        CodeRelationship(
            source_symbol="a",
            target_symbol="b",
            relationship_type=RelationshipType.CALL,
            source_file=str(file_path),
            target_file=None,
            source_line=2,
        ),
        CodeRelationship(
            source_symbol="b",
            target_symbol="c",
            relationship_type=RelationshipType.CALL,
            source_file=str(file_path),
            target_file=None,
            source_line=5,
        ),
    ]

    config = Config(data_dir=root / "data")
    store = DirIndexStore(index_db_path, config=config)
    store.initialize()
    store.add_file(
        name=file_path.name,
        full_path=file_path,
        content=content,
        language="python",
        symbols=symbols,
        relationships=relationships,
    )
    _compute_graph_neighbors(store)
    store.close()

    return mapper, project_root, file_path


def test_graph_neighbors_precomputed_two_hop(temp_paths: Path) -> None:
    mapper, project_root, file_path = _create_index_with_neighbors(temp_paths)
    index_db_path = mapper.source_to_index_db(project_root)

    conn = sqlite3.connect(str(index_db_path))
    conn.row_factory = sqlite3.Row
    try:
        rows = conn.execute(
            """
            SELECT s1.name AS source_name, s2.name AS neighbor_name, gn.relationship_depth
            FROM graph_neighbors gn
            JOIN symbols s1 ON s1.id = gn.source_symbol_id
            JOIN symbols s2 ON s2.id = gn.neighbor_symbol_id
            ORDER BY source_name, neighbor_name, relationship_depth
            """
        ).fetchall()
    finally:
        conn.close()

    triples = {(r["source_name"], r["neighbor_name"], int(r["relationship_depth"])) for r in rows}
    assert ("a", "b", 1) in triples
    assert ("a", "c", 2) in triples
    assert ("b", "c", 1) in triples
    assert ("c", "b", 1) in triples
    assert file_path.exists()


def test_graph_expander_returns_related_results_with_depth_metadata(temp_paths: Path) -> None:
    mapper, project_root, file_path = _create_index_with_neighbors(temp_paths)
    _ = project_root

    expander = GraphExpander(mapper, config=Config(data_dir=temp_paths / "data", graph_expansion_depth=2))
    base = SearchResult(
        path=str(file_path.resolve()),
        score=1.0,
        excerpt="",
        content=None,
        start_line=1,
        end_line=2,
        symbol_name="a",
        symbol_kind="function",
    )
    related = expander.expand([base], depth=2, max_expand=1, max_related=10)

    depth_by_symbol = {r.symbol_name: r.metadata.get("relationship_depth") for r in related}
    assert depth_by_symbol.get("b") == 1
    assert depth_by_symbol.get("c") == 2


def test_chain_search_populates_related_results_when_enabled(temp_paths: Path) -> None:
    mapper, project_root, file_path = _create_index_with_neighbors(temp_paths)
    _ = file_path

    registry = RegistryStore(db_path=temp_paths / "registry.db")
    registry.initialize()

    config = Config(
        data_dir=temp_paths / "data",
        enable_graph_expansion=True,
        graph_expansion_depth=2,
    )
    engine = ChainSearchEngine(registry, mapper, config=config)
    try:
        options = SearchOptions(depth=0, total_limit=10, enable_fuzzy=False)
        result = engine.search("b", project_root, options)

        assert result.results
        assert result.results[0].symbol_name == "a"

        depth_by_symbol = {r.symbol_name: r.metadata.get("relationship_depth") for r in result.related_results}
        assert depth_by_symbol.get("b") == 1
        assert depth_by_symbol.get("c") == 2
    finally:
        engine.close()


def test_chain_search_related_results_empty_when_disabled(temp_paths: Path) -> None:
    mapper, project_root, file_path = _create_index_with_neighbors(temp_paths)
    _ = file_path

    registry = RegistryStore(db_path=temp_paths / "registry.db")
    registry.initialize()

    config = Config(
        data_dir=temp_paths / "data",
        enable_graph_expansion=False,
    )
    engine = ChainSearchEngine(registry, mapper, config=config)
    try:
        options = SearchOptions(depth=0, total_limit=10, enable_fuzzy=False)
        result = engine.search("b", project_root, options)
        assert result.related_results == []
    finally:
        engine.close()