knowledge-base-framework/tests/test_kb_framework.py

#!/usr/bin/env python3
"""Tests for the importable `kb_framework` package.

The package is the contract both the stdio MCP server and the hosted connector
consume: build the graph in memory from a vault's Markdown, answer the query
tools over that graph, and read framework instructions. These tests pin the
graph shape against a golden fixture produced by the file-writing generator,
and the tool payloads against goldens captured from the stdio server, so the
two front ends cannot drift apart.
"""

from __future__ import annotations

import json
import sys
import unittest
from pathlib import Path

FRAMEWORK = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(FRAMEWORK))
FIXTURES = FRAMEWORK / "tests" / "fixtures"
VAULT = FIXTURES / "vault"

from kb_framework import graph as kbg  # noqa: E402
from kb_framework import query as kbq  # noqa: E402
from kb_framework import query_tools as kbt  # noqa: E402
from kb_framework import vault as kbv  # noqa: E402
from kb_framework import instructions as kbi  # noqa: E402


def golden_graph() -> dict:
    return json.loads((FIXTURES / "graph_golden.json").read_text(encoding="utf-8"))


def golden_stdio() -> dict:
    return json.loads((FIXTURES / "stdio_golden.json").read_text(encoding="utf-8"))


class TestBuildGraphInMemory(unittest.TestCase):
    def test_matches_golden_without_touching_disk(self) -> None:
        result = kbg.build_graph(VAULT)
        self.assertEqual(result.graph["nodes"], golden_graph()["nodes"])
        self.assertEqual(result.graph["edges"], golden_graph()["edges"])
        self.assertEqual(result.graph["version"], 1)
        self.assertTrue(result.graph["generated_at"])
        self.assertFalse((VAULT / "index").exists())

    def test_reports_no_problems_for_clean_fixture(self) -> None:
        result = kbg.build_graph(VAULT)
        self.assertEqual(result.unresolved, [])
        self.assertEqual(result.wiki_uses, [])
        self.assertEqual(kbg.validation_problems(result), [])

    def test_validation_reports_unresolved_and_bad_types(self) -> None:
        import tempfile

        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            (root / "a.md").write_text(
                "---\nid: a\ntype: descision\n---\n[x](missing){\"type\":\"bogus\"}\n[[legacy]]\n",
                encoding="utf-8",
            )
            result = kbg.build_graph(root)
            problems = "\n".join(kbg.validation_problems(result))
            self.assertIn("missing", problems)
            self.assertIn("descision", problems)
            self.assertIn("legacy", problems)
            self.assertFalse(result.ok)


class TestQueryFunctions(unittest.TestCase):
    def setUp(self) -> None:
        self.graph = golden_graph()
        self.stdio = golden_stdio()

    def strip(self, payload: dict) -> dict:
        payload = dict(payload)
        payload.pop("generated_at", None)
        return payload

    def test_summary_matches_stdio_golden(self) -> None:
        case = self.stdio["graph_summary"]
        payload = kbq.summary(self.graph, **case["args"])
        self.assertEqual(self.strip(payload), case["structuredContent"])

    def test_search_matches_stdio_golden(self) -> None:
        case = self.stdio["graph_search_nodes"]
        self.assertEqual(kbq.search_nodes(self.graph, case["args"]), case["structuredContent"])

    def test_neighbors_matches_stdio_golden(self) -> None:
        case = self.stdio["graph_neighbors"]
        self.assertEqual(kbq.neighbors(self.graph, **case["args"]), case["structuredContent"])

    def test_neighbors_unknown_id_raises(self) -> None:
        with self.assertRaises(ValueError):
            kbq.neighbors(self.graph, "nope")

    def test_node_path_resolves_id(self) -> None:
        self.assertEqual(kbq.node_path(self.graph, "area-ops"), "02_areas/ops.md")
        with self.assertRaises(ValueError):
            kbq.node_path(self.graph, "nope")


class TestVaultHelpers(unittest.TestCase):
    def test_resolve_in_root_rejects_escape_and_absolute(self) -> None:
        for bad in ("../x.md", "/etc/passwd", "", "   "):
            with self.subTest(bad=bad), self.assertRaises(ValueError):
                kbv.resolve_in_root(VAULT, bad)
        self.assertEqual(kbv.resolve_in_root(VAULT, "README.md"), VAULT / "README.md")

    def test_read_note_payload_text_and_binary(self) -> None:
        text = kbv.read_note_payload(VAULT, "02_areas/ops.md")
        self.assertEqual(text["binary"], False)
        self.assertIn("# Area Ops", text["content"])
        binary = kbv.read_note_payload(VAULT, "03_resources/demo/raw/paper.pdf")
        self.assertTrue(binary["binary"])
        self.assertEqual(set(binary), {"path", "binary", "bytes", "sha256", "hint"})
        with self.assertRaises(FileNotFoundError):
            kbv.read_note_payload(VAULT, "nope.md")


class TestQueryTools(unittest.TestCase):
    """The five read-only tools, as the stdio server and the connector expose them."""

    def setUp(self) -> None:
        self.graph = golden_graph()
        self.stdio = golden_stdio()

    def test_definitions_match_stdio(self) -> None:
        self.assertEqual(kbt.definitions(), self.stdio["tool_definitions"])

    def test_call_returns_stdio_result_shape(self) -> None:
        for key in ("graph_summary", "graph_search_nodes", "graph_neighbors", "read_note_by_id", "read_note_by_path"):
            case = self.stdio[key]
            with self.subTest(tool=key):
                result = kbt.call(case["tool"], case["args"], graph=self.graph, root=VAULT)
                sc = dict(result["structuredContent"])
                sc.pop("generated_at", None)
                self.assertEqual(sc, case["structuredContent"])
                self.assertEqual(result["content"][0]["type"], "text")
                self.assertEqual(json.loads(result["content"][0]["text"]).get("path", sc.get("path")), sc.get("path"))

    def test_read_note_requires_exactly_one_selector(self) -> None:
        for args in ({}, {"id": "area-ops", "path": "README.md"}):
            with self.subTest(args=args), self.assertRaises(ValueError):
                kbt.call("read_note", args, graph=self.graph, root=VAULT)

    def test_unknown_tool_raises(self) -> None:
        with self.assertRaises(ValueError):
            kbt.call("write_note", {}, graph=self.graph, root=VAULT)


class TestInstructions(unittest.TestCase):
    def test_reads_framework_documents(self) -> None:
        payload = kbi.read_instructions()
        self.assertEqual([d["path"] for d in payload["documents"]], ["AGENTS.md", "AGENT_QUERY_GUIDE.md"])
        self.assertIn("AGENT_QUERY_GUIDE", payload["documents"][0]["content"])
        self.assertTrue(Path(payload["framework_root"]).is_dir())


if __name__ == "__main__":
    unittest.main()