grammar-inference-engine/tests/test_grammar_index.py

"""Tests for bex.grammar_index — runtime grammar lookup."""

import os
import tempfile

import pytest
import yaml

from bex.grammar_index import GrammarIndex, load_grammar_index


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

def _make_yml(tmp_path, entries):
    """Write a grammars.yml with the given entries grouped by module."""
    dervish = tmp_path / ".dervish"
    dervish.mkdir()
    data = {"test_module": entries}
    with open(dervish / "grammars.yml", "w") as f:
        yaml.dump(data, f, default_flow_style=False)


SAMPLE_ENTRIES = [
    {
        "package": "src/services [return]",
        "methods": 10,
        "grammar": "return.ok+?.error?",
        "score": 0.8,
        "algorithm": "CRX",
        "mdl": 12.0,
    },
    {
        "package": "src/services [if]",
        "methods": 5,
        "grammar": "if.cond+.then+.else?",
        "score": 0.6,
        "algorithm": "CRX",
        "mdl": 15.0,
    },
    {
        "package": "src/utils [return]",
        "methods": 8,
        "grammar": "return.value+",
        "score": 1.0,
        "algorithm": "CRX",
        "mdl": 5.0,
    },
]


# ---------------------------------------------------------------------------
# GrammarIndex construction
# ---------------------------------------------------------------------------

class TestGrammarIndex:

    def test_parse_label_with_symbol(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        pkgs = idx.packages()
        assert "src/services" in pkgs
        assert "src/utils" in pkgs

    def test_parse_label_without_symbol(self, tmp_path):
        entries = [{"package": "src/root", "grammar": "a.b", "score": 0.5}]
        idx = GrammarIndex(str(tmp_path), entries)
        assert "src/root" in idx.packages()

    def test_sorted_by_score_desc(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        entries = idx.get_package("src/services")
        scores = [s for _, _, s, _ in entries]
        assert scores == sorted(scores, reverse=True)

    def test_repr(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        assert "3 entries" in repr(idx)
        assert "2 packages" in repr(idx)


# ---------------------------------------------------------------------------
# get()
# ---------------------------------------------------------------------------

class TestGet:

    def test_get_best_no_context(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        g = idx.get("src/services/main.kt")
        assert g == "return.ok+?.error?"

    def test_get_by_context_symbol(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        g = idx.get("src/services/main.kt", context_symbol="if")
        assert g == "if.cond+.then+.else?"

    def test_get_missing_context_falls_back_to_best(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        g = idx.get("src/services/main.kt", context_symbol="nonexistent")
        assert g == "return.ok+?.error?"

    def test_get_unknown_package_returns_none(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        g = idx.get("src/unknown/file.kt")
        assert g is None

    def test_get_absolute_path(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        abs_path = str(tmp_path / "src" / "services" / "main.kt")
        g = idx.get(abs_path, context_symbol="return")
        assert g == "return.ok+?.error?"


# ---------------------------------------------------------------------------
# get_package()
# ---------------------------------------------------------------------------

class TestGetPackage:

    def test_returns_all_entries(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        entries = idx.get_package("src/services/main.kt")
        assert len(entries) == 2

    def test_unknown_package_returns_empty(self, tmp_path):
        idx = GrammarIndex(str(tmp_path), SAMPLE_ENTRIES)
        entries = idx.get_package("src/unknown/file.kt")
        assert entries == []


# ---------------------------------------------------------------------------
# load_grammar_index()
# ---------------------------------------------------------------------------

class TestLoadGrammarIndex:

    def test_load_existing(self, tmp_path):
        _make_yml(tmp_path, SAMPLE_ENTRIES)
        idx = load_grammar_index(str(tmp_path))
        assert len(idx.get_all()) == 3
        assert "src/services" in idx.packages()

    def test_load_missing_file(self, tmp_path):
        idx = load_grammar_index(str(tmp_path))
        assert len(idx.get_all()) == 0
        assert idx.packages() == []

    def test_roundtrip_persisted_ragsak(self):
        """Integration test: grammars.yml generated by analyze_directory."""
        ragsak_yml = "/home/tobi/Desktop/kesai/RAGSAK/.dervish/grammars.yml"
        if not os.path.exists(ragsak_yml):
            pytest.skip("RAGSAK grammars.yml not generated yet")
        idx = load_grammar_index("/home/tobi/Desktop/kesai/RAGSAK")
        assert len(idx.get_all()) > 0
        # Should be able to resolve a file in a package with a grammar
        g = idx.get(
            "/home/tobi/Desktop/kesai/RAGSAK/agents/capability/support/src/main/kotlin/eu/corentic/springrag/agent/capability/support/AgentCapabilityConfiguration.kt",
            context_symbol="id",
        )
        assert g is not None
        assert "id" in g