"""Tests for the M14 AI assistant: provider selection, grounding, parsing (no network).""" import unittest from unittest import mock from rigdoctor.core import ai, ai_knowledge class KnowledgeTests(unittest.TestCase): def test_matches_xid_and_smart(self): facts = ai_knowledge.relevant("Kernel: NVRM: Xid 79: GPU has fallen off the bus") self.assertTrue(any("fallen off the bus" in f for f in facts)) def test_matches_smart_pending(self): facts = ai_knowledge.relevant("SMART 197 Current_Pending_Sector = 8") self.assertTrue(any("Pending Sector" in f for f in facts)) def test_no_match_returns_empty(self): self.assertEqual(ai_knowledge.relevant("everything is fine"), []) class ConfigStateTests(unittest.TestCase): def _cfg(self, **over): base = {"ai_provider": "", "ai_model": "", "ai_endpoint": "http://localhost:11434"} base.update(over) return base def test_unconfigured_by_default(self): with mock.patch.object(ai.config, "load_config", return_value=self._cfg()): self.assertFalse(ai.is_configured()) def test_ollama_needs_model(self): with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="ollama")): self.assertFalse(ai.is_configured()) with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="ollama", ai_model="llama3.1")): self.assertTrue(ai.is_configured()) def test_claude_needs_key(self): with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")), \ mock.patch.object(ai.config, "load_ai_key", return_value=None): self.assertFalse(ai.is_configured()) with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")), \ mock.patch.object(ai.config, "load_ai_key", return_value="sk-ant-x"): self.assertTrue(ai.is_configured()) def test_claude_default_model(self): with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")): self.assertEqual(ai.model(), ai.CLAUDE_DEFAULT_MODEL) class PromptTests(unittest.TestCase): def test_build_prompt_includes_facts_and_findings(self): prompt = ai.build_prompt("Xid 79: GPU has fallen off the bus") self.assertIn("Reference facts", prompt) self.assertIn("Collected findings", prompt) self.assertIn("fallen off the bus", prompt) def test_format_findings(self): class F: severity, category, title, detail = "warn", "GPU", "Hot", "92C" text = ai.format_findings([F()]) self.assertIn("[WARN] GPU: Hot — 92C", text) def test_appid_glossary_resolves_known_ids(self): from rigdoctor.core import steam with mock.patch.object(steam, "appid_names", return_value={"2694490": "Path of Exile 2"}): glossary = ai.appid_glossary("Steam log: removed AppID 2694490 ... pid 130544") self.assertIn("2694490 = Path of Exile 2", glossary) def test_appid_glossary_ignores_unknown_ids(self): from rigdoctor.core import steam with mock.patch.object(steam, "appid_names", return_value={"570": "Dota 2"}): self.assertEqual(ai.appid_glossary("pid 130544 used 8192 MiB"), "") # not in library def test_build_prompt_includes_glossary(self): from rigdoctor.core import steam with mock.patch.object(steam, "appid_names", return_value={"2694490": "Path of Exile 2"}): prompt = ai.build_prompt("AppID 2694490 launched") self.assertIn("Path of Exile 2", prompt) class ExplainTests(unittest.TestCase): def _cfg(self, **over): base = {"ai_provider": "", "ai_model": "", "ai_endpoint": "http://localhost:11434"} base.update(over) return base def test_no_provider(self): with mock.patch.object(ai.config, "load_config", return_value=self._cfg()): ok, msg = ai.explain("x") self.assertFalse(ok) self.assertIn("No AI provider", msg) def test_ollama_parses_response(self): with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="ollama", ai_model="llama3.1")), \ mock.patch.object(ai, "_post", return_value={"response": "It's the PSU."}) as post: ok, msg = ai.explain("Xid 79") self.assertTrue(ok) self.assertEqual(msg, "It's the PSU.") self.assertIn("/api/generate", post.call_args[0][0]) def test_claude_parses_content_blocks(self): with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")), \ mock.patch.object(ai.config, "load_ai_key", return_value="sk-ant-x"), \ mock.patch.object(ai, "_post", return_value={"content": [ {"type": "text", "text": "Likely a failing disk."}]}) as post: ok, msg = ai.explain("SMART 197") self.assertTrue(ok) self.assertEqual(msg, "Likely a failing disk.") headers = post.call_args[0][2] self.assertEqual(headers["anthropic-version"], ai.ANTHROPIC_VERSION) self.assertEqual(headers["x-api-key"], "sk-ant-x") class _FakeResp: """A context-managed iterable of byte lines, like urlopen() returns.""" def __init__(self, lines): self._lines = [l.encode("utf-8") for l in lines] def __enter__(self): return iter(self._lines) def __exit__(self, *a): return False class StreamTests(unittest.TestCase): def _cfg(self, **over): base = {"ai_provider": "", "ai_model": "", "ai_endpoint": "http://localhost:11434"} base.update(over) return base def test_ollama_stream_accumulates_and_callbacks(self): lines = ['{"response": "It is ", "done": false}', '{"response": "the PSU.", "done": false}', '{"response": "", "done": true}'] chunks = [] with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="ollama", ai_model="qwen2.5:7b")), \ mock.patch.object(ai, "_stream_request", return_value=_FakeResp(lines)): ok, full = ai.explain_stream("Xid 79", on_chunk=chunks.append) self.assertTrue(ok) self.assertEqual(full, "It is the PSU.") self.assertEqual(chunks, ["It is ", "the PSU."]) def test_claude_stream_parses_sse(self): lines = [ 'event: content_block_delta', 'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"Failing "}}', 'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"disk."}}', 'data: {"type":"message_stop"}', ] chunks = [] with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")), \ mock.patch.object(ai.config, "load_ai_key", return_value="sk-ant-x"), \ mock.patch.object(ai, "_stream_request", return_value=_FakeResp(lines)): ok, full = ai.explain_stream("SMART 197", on_chunk=chunks.append) self.assertTrue(ok) self.assertEqual(full, "Failing disk.") self.assertEqual(chunks, ["Failing ", "disk."]) if __name__ == "__main__": unittest.main()