rigdoctor/tests/test_ai.py

"""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()