test: add 31 tests for RAG engine and rule plugin manager

Author: mnbplus

tests/test_rag.py

@@ -0,0 +1,201 @@
+"""Tests for pyaegis.rag -- local RAG engine."""
+from __future__ import annotations
+
+import textwrap
+from pathlib import Path
+
+from pyaegis.rag import CodeRAG, _chunk_file, _cosine_similarity
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _write_py(tmp_path: Path, name: str, src: str) -> Path:
+    p = tmp_path / name
+    p.write_text(textwrap.dedent(src), encoding="utf-8")
+    return p
+
+
+# ---------------------------------------------------------------------------
+# _chunk_file
+# ---------------------------------------------------------------------------
+
+
+def test_chunk_file_functions(tmp_path):
+    f = _write_py(
+        tmp_path,
+        "sample.py",
+        (
+            "def foo(x):\n"
+            "    # Foo does stuff.\n"
+            "    return x + 1\n"
+            "\n"
+            "def bar():\n"
+            "    pass\n"
+        ),
+    )
+    chunks = _chunk_file(str(f))
+    names = [c.name for c in chunks]
+    assert "foo" in names
+    assert "bar" in names
+    for c in chunks:
+        assert c.kind == "function"
+        assert c.file_path == str(f)
+        assert c.chunk_id
+
+
+def test_chunk_file_class(tmp_path):
+    f = _write_py(
+        tmp_path,
+        "cls.py",
+        (
+            "class MyClass:\n"
+            "    # A class.\n"
+            "    def method(self):\n"
+            "        pass\n"
+        ),
+    )
+    chunks = _chunk_file(str(f))
+    kinds = {c.kind for c in chunks}
+    assert "class" in kinds
+
+
+def test_chunk_file_nonexistent():
+    chunks = _chunk_file("/does/not/exist.py")
+    assert chunks == []
+
+
+def test_chunk_file_syntax_error(tmp_path):
+    f = _write_py(tmp_path, "bad.py", "def foo(:\n    pass\n")
+    chunks = _chunk_file(str(f))
+    assert isinstance(chunks, list)
+
+
+# ---------------------------------------------------------------------------
+# _cosine_similarity
+# ---------------------------------------------------------------------------
+
+
+def test_cosine_similarity_identical():
+    v = [1.0, 0.0, 0.0]
+    assert abs(_cosine_similarity(v, v) - 1.0) < 1e-6
+
+
+def test_cosine_similarity_orthogonal():
+    v1 = [1.0, 0.0]
+    v2 = [0.0, 1.0]
+    assert abs(_cosine_similarity(v1, v2)) < 1e-6
+
+
+def test_cosine_similarity_zero_vector():
+    v = [0.0, 0.0, 0.0]
+    assert _cosine_similarity(v, v) == 0.0
+
+
+# ---------------------------------------------------------------------------
+# CodeRAG -- basic index + search
+# ---------------------------------------------------------------------------
+
+
+def test_rag_index_and_search(tmp_path):
+    f = _write_py(
+        tmp_path,
+        "auth.py",
+        (
+            "def authenticate_user(username, password):\n"
+            "    # Check user credentials against the database.\n"
+            "    return check_db(username, password)\n"
+            "\n"
+            "def logout(session):\n"
+            "    # Invalidate the user session.\n"
+            "    session.clear()\n"
+        ),
+    )
+    rag = CodeRAG(db_path=":memory:")
+    n = rag.index_file(str(f))
+    assert n >= 2
+
+    results = rag.search("user authentication credentials", top_k=3)
+    assert len(results) > 0
+    names = [r.chunk.name for r in results]
+    assert "authenticate_user" in names
+    rag.close()
+
+
+def test_rag_index_directory(tmp_path):
+    _write_py(tmp_path, "a.py", "def alpha(): pass\n")
+    _write_py(tmp_path, "b.py", "def beta(): pass\n")
+    sub = tmp_path / "sub"
+    sub.mkdir()
+    _write_py(sub, "c.py", "def gamma(): pass\n")
+
+    rag = CodeRAG(db_path=":memory:")
+    files, chunks = rag.index_directory(str(tmp_path))
+    assert files >= 2
+    assert chunks >= 2
+    stats = rag.stats()
+    assert stats["indexed_files"] >= 2
+    rag.close()
+
+
+def test_rag_no_reindex_unchanged(tmp_path):
+    f = _write_py(tmp_path, "stable.py", "def stable(): pass\n")
+    db = tmp_path / "rag.sqlite"
+    rag = CodeRAG(db_path=str(db))
+    n1 = rag.index_file(str(f))
+    n2 = rag.index_file(str(f))
+    assert n1 > 0
+    assert n2 == 0
+    rag.close()
+
+
+def test_rag_force_reindex(tmp_path):
+    f = _write_py(tmp_path, "x.py", "def x(): pass\n")
+    db = tmp_path / "rag.sqlite"
+    rag = CodeRAG(db_path=str(db))
+    rag.index_file(str(f))
+    n = rag.index_file(str(f), force=True)
+    assert n > 0
+    rag.close()
+
+
+def test_rag_kind_filter(tmp_path):
+    f = _write_py(
+        tmp_path,
+        "mixed.py",
+        ("class Foo:\n" "    pass\n" "\n" "def bar():\n" "    pass\n"),
+    )
+    rag = CodeRAG(db_path=":memory:")
+    rag.index_file(str(f))
+    results = rag.search("foo bar", top_k=10, kind_filter="function")
+    for r in results:
+        assert r.chunk.kind == "function"
+    rag.close()
+
+
+def test_rag_build_context(tmp_path):
+    f = _write_py(tmp_path, "ctx.py", "def hello(): pass\n")
+    rag = CodeRAG(db_path=":memory:")
+    rag.index_file(str(f))
+    results = rag.search("hello", top_k=1)
+    ctx = rag.build_context(results, max_chars=5000)
+    assert "hello" in ctx
+    rag.close()
+
+
+def test_rag_stats_empty():
+    rag = CodeRAG(db_path=":memory:")
+    s = rag.stats()
+    assert s["total_chunks"] == 0
+    assert s["indexed_files"] == 0
+    rag.close()
+
+
+def test_rag_context_manager(tmp_path):
+    f = _write_py(tmp_path, "cm.py", "def cm(): pass\n")
+    with CodeRAG(db_path=":memory:") as rag:
+        rag.index_file(str(f))
+        results = rag.search("cm")
+        assert isinstance(results, list)

tests/test_rule_plugins.py

@@ -0,0 +1,177 @@
+"""Tests for pyaegis.rule_plugins — community rule pack manager."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import yaml
+
+from pyaegis.rule_plugins import RulePluginManager
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def mgr(tmp_path):
+    return RulePluginManager(rules_dir=str(tmp_path / "rules"))
+
+
+@pytest.fixture
+def sample_pack(tmp_path):
+    """Write a valid YAML rule pack to disk and return its path."""
+    data = {
+        "inputs": ["environ.get", "getenv"],
+        "sinks": ["ldap.search"],
+        "sanitizers": ["escape"],
+        "conditional_sinks": [],
+        "source_decorators": [],
+    }
+    p = tmp_path / "sample_pack.yml"
+    p.write_text(yaml.safe_dump(data), encoding="utf-8")
+    return str(p)
+
+
+# ---------------------------------------------------------------------------
+# Install
+# ---------------------------------------------------------------------------
+
+
+def test_install_local(mgr, sample_pack):
+    name = mgr.install(sample_pack, name="test-pack")
+    assert name == "test-pack"
+    installed = mgr.list_installed()
+    assert len(installed) == 1
+    assert installed[0]["name"] == "test-pack"
+
+
+def test_install_auto_name(mgr, sample_pack):
+    name = mgr.install(sample_pack)  # name inferred from filename
+    assert name == "sample_pack"
+
+
+def test_install_duplicate_raises(mgr, sample_pack):
+    mgr.install(sample_pack, name="dup")
+    with pytest.raises(ValueError, match="already installed"):
+        mgr.install(sample_pack, name="dup")
+
+
+def test_install_duplicate_force(mgr, sample_pack):
+    mgr.install(sample_pack, name="dup")
+    name = mgr.install(sample_pack, name="dup", force=True)
+    assert name == "dup"
+    assert len(mgr.list_installed()) == 1
+
+
+def test_install_invalid_yaml(mgr, tmp_path):
+    bad = tmp_path / "bad.yml"
+    bad.write_text("[unclosed bracket\n", encoding="utf-8")
+    with pytest.raises(ValueError, match="Invalid YAML"):
+        mgr.install(str(bad), name="bad")
+
+
+def test_install_non_mapping_yaml(mgr, tmp_path):
+    p = tmp_path / "list.yml"
+    p.write_text("- item1\n- item2\n", encoding="utf-8")
+    with pytest.raises(ValueError, match="must be a YAML mapping"):
+        mgr.install(str(p), name="list")
+
+
+# ---------------------------------------------------------------------------
+# Remove
+# ---------------------------------------------------------------------------
+
+
+def test_remove_installed(mgr, sample_pack):
+    mgr.install(sample_pack, name="to-remove")
+    result = mgr.remove("to-remove")
+    assert result is True
+    assert mgr.list_installed() == []
+
+
+def test_remove_nonexistent(mgr):
+    result = mgr.remove("ghost")
+    assert result is False
+
+
+# ---------------------------------------------------------------------------
+# List
+# ---------------------------------------------------------------------------
+
+
+def test_list_empty(mgr):
+    assert mgr.list_installed() == []
+
+
+def test_list_multiple(mgr, tmp_path, sample_pack):
+    pack2 = tmp_path / "pack2.yml"
+    pack2.write_text(yaml.safe_dump({"sinks": ["eval"]}), encoding="utf-8")
+    mgr.install(sample_pack, name="p1")
+    mgr.install(str(pack2), name="p2")
+    names = {p["name"] for p in mgr.list_installed()}
+    assert names == {"p1", "p2"}
+
+
+# ---------------------------------------------------------------------------
+# Merge
+# ---------------------------------------------------------------------------
+
+
+def test_merged_rules_includes_plugin(mgr, sample_pack):
+    mgr.install(sample_pack, name="extra")
+    merged = mgr.merged_rules(include_builtin=False)
+    assert "environ.get" in merged["inputs"]
+    assert "ldap.search" in merged["sinks"]
+
+
+def test_merged_rules_deduplicates(mgr, tmp_path):
+    p1 = tmp_path / "p1.yml"
+    p2 = tmp_path / "p2.yml"
+    p1.write_text(yaml.safe_dump({"inputs": ["os.environ"]}), encoding="utf-8")
+    p2.write_text(
+        yaml.safe_dump({"inputs": ["os.environ", "getenv"]}), encoding="utf-8"
+    )
+    mgr.install(str(p1), name="p1")
+    mgr.install(str(p2), name="p2")
+    merged = mgr.merged_rules(include_builtin=False)
+    assert merged["inputs"].count("os.environ") == 1
+
+
+def test_merged_rules_path_creates_file(mgr, sample_pack):
+    mgr.install(sample_pack, name="mp")
+    path = mgr.merged_rules_path(include_builtin=False)
+    assert Path(path).exists()
+    content = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
+    assert "sinks" in content
+    Path(path).unlink()  # cleanup
+
+
+# ---------------------------------------------------------------------------
+# has_plugins
+# ---------------------------------------------------------------------------
+
+
+def test_has_plugins_false(mgr):
+    assert mgr.has_plugins() is False
+
+
+def test_has_plugins_true(mgr, sample_pack):
+    mgr.install(sample_pack, name="any")
+    assert mgr.has_plugins() is True
+
+
+# ---------------------------------------------------------------------------
+# Index persistence
+# ---------------------------------------------------------------------------
+
+
+def test_index_persists_across_instances(tmp_path, sample_pack):
+    rules_dir = str(tmp_path / "rules")
+    mgr1 = RulePluginManager(rules_dir=rules_dir)
+    mgr1.install(sample_pack, name="persist")
+
+    mgr2 = RulePluginManager(rules_dir=rules_dir)
+    names = {p["name"] for p in mgr2.list_installed()}
+    assert "persist" in names