280 lines
9.5 KiB
Python
280 lines
9.5 KiB
Python
"""Tests for runbook_store — no network calls, ChromaDB mocked."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from tai.runbook_store import (
|
|
RunbookChunk,
|
|
RunbookMeta,
|
|
RunbookStore,
|
|
_build_embed_text,
|
|
_parse_frontmatter,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _parse_frontmatter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_parse_frontmatter_extracts_service() -> None:
|
|
text = "---\nservice: nginx\n---\n## Body\nsome content"
|
|
meta, body = _parse_frontmatter(text)
|
|
assert meta.service == "nginx"
|
|
assert "## Body" in body
|
|
|
|
|
|
def test_parse_frontmatter_extracts_tags_as_list() -> None:
|
|
text = "---\ntags: nginx, web, http\n---\nbody"
|
|
meta, body = _parse_frontmatter(text)
|
|
assert meta.tags == ["nginx", "web", "http"]
|
|
|
|
|
|
def test_parse_frontmatter_extracts_symptoms_as_list() -> None:
|
|
text = "---\nsymptoms: 502 Bad Gateway, upstream refused\n---\nbody"
|
|
meta, body = _parse_frontmatter(text)
|
|
assert meta.symptoms == ["502 Bad Gateway", "upstream refused"]
|
|
|
|
|
|
def test_parse_frontmatter_returns_empty_meta_when_missing() -> None:
|
|
text = "# Just a heading\nno frontmatter here"
|
|
meta, body = _parse_frontmatter(text)
|
|
assert meta.service == ""
|
|
assert meta.tags == []
|
|
assert meta.symptoms == []
|
|
assert "Just a heading" in body
|
|
|
|
|
|
def test_parse_frontmatter_body_strips_delimiter() -> None:
|
|
text = "---\nservice: ssh\n---\nBody starts here."
|
|
_, body = _parse_frontmatter(text)
|
|
assert body.strip() == "Body starts here."
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _build_embed_text
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_build_embed_text_includes_title_and_service() -> None:
|
|
meta = RunbookMeta(service="nginx", symptoms=["502"], tags=["web"])
|
|
result = _build_embed_text("nginx", meta, "body content")
|
|
assert "title: nginx" in result
|
|
assert "service: nginx" in result
|
|
|
|
|
|
def test_build_embed_text_includes_symptoms_and_tags() -> None:
|
|
meta = RunbookMeta(service="nginx", symptoms=["502 Bad Gateway"], tags=["web", "http"])
|
|
result = _build_embed_text("nginx", meta, "body")
|
|
assert "502 Bad Gateway" in result
|
|
assert "web" in result
|
|
|
|
|
|
def test_build_embed_text_includes_body_excerpt() -> None:
|
|
meta = RunbookMeta()
|
|
result = _build_embed_text("disk", meta, "check df -h output")
|
|
assert "check df -h output" in result
|
|
|
|
|
|
def test_build_embed_text_truncates_long_body() -> None:
|
|
meta = RunbookMeta()
|
|
long_body = "x" * 2000
|
|
result = _build_embed_text("disk", meta, long_body)
|
|
# Body excerpt is capped at 800 chars
|
|
assert len(result) < 1500
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RunbookStore — unit tests using tmp_path and mocked chromadb
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_chromadb_mock() -> MagicMock:
|
|
"""Return a chromadb mock that satisfies RunbookStore internals."""
|
|
collection = MagicMock()
|
|
collection.count.return_value = 0
|
|
client = MagicMock()
|
|
client.get_or_create_collection.return_value = collection
|
|
chroma_mod = MagicMock()
|
|
chroma_mod.PersistentClient.return_value = client
|
|
chroma_mod.HttpClient.return_value = client
|
|
return chroma_mod
|
|
|
|
|
|
def _make_ai_mock(embedding: list[float] | None = None) -> MagicMock:
|
|
ai = MagicMock()
|
|
ai.embed.return_value = embedding or [0.1, 0.2, 0.3]
|
|
return ai
|
|
|
|
|
|
def test_runbook_store_sync_returns_count(tmp_path: Path) -> None:
|
|
(tmp_path / "nginx.md").write_text(
|
|
"---\nservice: nginx\ntags: web\nsymptoms: 502\n---\n## Body\ncontent"
|
|
)
|
|
(tmp_path / "ssh.md").write_text(
|
|
"---\nservice: ssh\ntags: ssh\nsymptoms: refused\n---\n## Body\ncontent"
|
|
)
|
|
|
|
chroma_mock = _make_chromadb_mock()
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
count = store.sync(tmp_path, ai)
|
|
|
|
assert count == 2
|
|
|
|
|
|
def test_runbook_store_sync_calls_upsert(tmp_path: Path) -> None:
|
|
(tmp_path / "nginx.md").write_text("---\nservice: nginx\n---\nbody")
|
|
|
|
chroma_mock = _make_chromadb_mock()
|
|
collection = chroma_mock.PersistentClient.return_value.get_or_create_collection.return_value
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
store.sync(tmp_path, ai)
|
|
|
|
collection.upsert.assert_called_once()
|
|
call_kwargs = collection.upsert.call_args.kwargs
|
|
assert "nginx" in call_kwargs["ids"]
|
|
|
|
|
|
def test_runbook_store_sync_empty_dir_returns_zero(tmp_path: Path) -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
count = store.sync(tmp_path, ai)
|
|
|
|
assert count == 0
|
|
|
|
|
|
def test_runbook_store_sync_missing_dir_raises(tmp_path: Path) -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
with pytest.raises(FileNotFoundError):
|
|
store.sync(tmp_path / "nonexistent", ai)
|
|
|
|
|
|
def test_runbook_store_query_returns_empty_when_no_docs(tmp_path: Path) -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
# collection.count() returns 0 by default in our mock
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
results = store.query("disk full", ai)
|
|
|
|
assert results == []
|
|
|
|
|
|
def test_runbook_store_query_returns_runbook_chunks(tmp_path: Path) -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
collection = chroma_mock.PersistentClient.return_value.get_or_create_collection.return_value
|
|
collection.count.return_value = 2
|
|
collection.query.return_value = {
|
|
"documents": [["## Body\ncheck df -h"]],
|
|
"metadatas": [
|
|
[{"title": "disk", "service": "disk", "tags": "disk, storage", "symptoms": "full"}]
|
|
],
|
|
}
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
results = store.query("disk is full", ai)
|
|
|
|
assert len(results) == 1
|
|
assert isinstance(results[0], RunbookChunk)
|
|
assert results[0].title == "disk"
|
|
assert results[0].service == "disk"
|
|
assert "disk" in results[0].tags
|
|
assert "df -h" in results[0].content
|
|
|
|
|
|
def test_runbook_store_list_indexed_returns_metadata(tmp_path: Path) -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
collection = chroma_mock.PersistentClient.return_value.get_or_create_collection.return_value
|
|
collection.count.return_value = 1
|
|
collection.get.return_value = {
|
|
"metadatas": [{"title": "nginx", "service": "nginx", "tags": "web", "symptoms": "502"}]
|
|
}
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
entries = store.list_indexed()
|
|
|
|
assert len(entries) == 1
|
|
assert entries[0]["title"] == "nginx"
|
|
|
|
|
|
def test_runbook_store_count_delegates_to_collection(tmp_path: Path) -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
collection = chroma_mock.PersistentClient.return_value.get_or_create_collection.return_value
|
|
collection.count.return_value = 5
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
assert store.count() == 5
|
|
|
|
|
|
def test_runbook_store_sync_single_upserts_one(tmp_path: Path) -> None:
|
|
runbook = tmp_path / "nginx.md"
|
|
runbook.write_text("---\nservice: nginx\ntags: web\n---\nbody text")
|
|
|
|
chroma_mock = _make_chromadb_mock()
|
|
collection = chroma_mock.PersistentClient.return_value.get_or_create_collection.return_value
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
store.sync_single(runbook, ai)
|
|
|
|
collection.upsert.assert_called_once()
|
|
call_kwargs = collection.upsert.call_args.kwargs
|
|
assert call_kwargs["ids"] == ["nginx"]
|
|
|
|
|
|
def test_runbook_store_sync_single_missing_file_raises(tmp_path: Path) -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
ai = _make_ai_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore(tmp_path / "store")
|
|
with pytest.raises(FileNotFoundError):
|
|
store.sync_single(tmp_path / "missing.md", ai)
|
|
|
|
|
|
def test_runbook_store_remote_url_uses_http_client() -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore("https://chroma.example.com:8443")
|
|
assert store.count() == 0
|
|
|
|
chroma_mock.HttpClient.assert_called_once_with(host="chroma.example.com", port=8443, ssl=True)
|
|
|
|
|
|
def test_runbook_store_remote_url_uses_http_client_with_basic_auth() -> None:
|
|
chroma_mock = _make_chromadb_mock()
|
|
|
|
with patch.dict("sys.modules", {"chromadb": chroma_mock}):
|
|
store = RunbookStore("https://chroma.example.com:8443", username="tai", password="secret")
|
|
assert store.count() == 0
|
|
|
|
args = chroma_mock.HttpClient.call_args.kwargs
|
|
assert args["host"] == "chroma.example.com"
|
|
assert args["port"] == 8443
|
|
assert args["ssl"] is True
|
|
assert "headers" in args
|
|
assert str(args["headers"].get("Authorization", "")).startswith("Basic ")
|