from __future__ import annotations import importlib.util import os from pathlib import Path from unittest import TestCase, mock BOT_PATH = Path(__file__).resolve().parents[1] / "atlasbot" / "bot.py" def load_bot_module(): env = { "BOT_USER": "atlas-smart", "BOT_PASS": "smart-pass", "BOT_USER_QUICK": "atlas-quick", "BOT_PASS_QUICK": "quick-pass", "BOT_USER_SMART": "atlas-smart", "BOT_PASS_SMART": "smart-pass", "BOT_USER_GENIUS": "atlas-genius", "BOT_PASS_GENIUS": "genius-pass", "OLLAMA_URL": "http://ollama.invalid", "OLLAMA_MODEL": "base-model", "ATLASBOT_MODEL_FAST": "fast-model", "ATLASBOT_MODEL_SMART": "smart-model", "ATLASBOT_MODEL_GENIUS": "genius-model", "ATLASBOT_QUICK_TIME_BUDGET_SEC": "15", "ATLASBOT_SMART_TIME_BUDGET_SEC": "45", "ATLASBOT_GENIUS_TIME_BUDGET_SEC": "180", "KB_DIR": "", "VM_URL": "http://vm.invalid", "ARIADNE_STATE_URL": "", "ARIADNE_STATE_TOKEN": "", } with mock.patch.dict(os.environ, env, clear=False): spec = importlib.util.spec_from_file_location("atlasbot_bot", BOT_PATH) module = importlib.util.module_from_spec(spec) assert spec.loader is not None spec.loader.exec_module(module) return module class AtlasbotModeTests(TestCase): def setUp(self): self.bot = load_bot_module() def test_bot_accounts_include_genius_mode(self): accounts = self.bot._bot_accounts() by_user = {account["user"]: account["mode"] for account in accounts} self.assertEqual(by_user["atlas-quick"], "fast") self.assertEqual(by_user["atlas-smart"], "smart") self.assertEqual(by_user["atlas-genius"], "genius") def test_objective_cluster_question_uses_fact_pack_without_llm(self): fact_lines = [ "hottest_cpu: longhorn-system (6.69)", "hottest_ram: longhorn-system (36.05 GB)", ] with ( mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines), mock.patch.object(self.bot, "_ollama_call", side_effect=AssertionError("LLM should not be called")), ): reply = self.bot.open_ended_answer( "what is the hottest cpu node in titan lab currently?", inventory=[], snapshot=None, workloads=[], history_lines=[], mode="smart", allow_tools=True, ) self.assertIn("longhorn-system", reply) self.assertIn("Confidence:", reply) def test_subjective_genius_answer_uses_genius_model(self): fact_lines = [ "hottest_cpu: longhorn-system (6.69)", "worker_nodes: titan-01, titan-02, titan-03", ] captured: dict[str, object] = {} def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None): captured["model"] = model captured["timeout"] = timeout captured["context"] = context return "The worker spread stands out because Titan keeps meaningful capacity on the same cluster. Confidence: high" with ( mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines), mock.patch.object(self.bot, "_ollama_call", side_effect=fake_ollama_call), ): reply = self.bot.open_ended_answer( "what stands out about titan lab?", inventory=[], snapshot=None, workloads=[], history_lines=[], mode="genius", allow_tools=True, context='Cluster snapshot (JSON): {"injected":true}', ) self.assertIn("The worker spread stands out", reply) self.assertEqual(captured["model"], "genius-model") self.assertLessEqual(float(captured["timeout"]), 180.0) self.assertIn('Cluster snapshot (JSON): {"injected":true}', str(captured["context"])) def test_mode_timeouts_stay_within_budgets(self): fact_lines = [ "hottest_cpu: longhorn-system (6.69)", "worker_nodes: titan-01, titan-02, titan-03", ] seen: list[tuple[str, float]] = [] def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None): seen.append((str(model), float(timeout or 0))) return "Atlas has a clear standout because the worker spread is healthy. Confidence: high" with ( mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines), mock.patch.object(self.bot, "_ollama_call", side_effect=fake_ollama_call), ): for mode in ("fast", "smart", "genius"): reply = self.bot.open_ended_answer( "what stands out about titan lab?", inventory=[], snapshot=None, workloads=[], history_lines=[], mode=mode, allow_tools=True, ) self.assertIn("Confidence:", reply) self.assertEqual([model for model, _ in seen], ["fast-model", "smart-model", "genius-model"]) self.assertLessEqual(seen[0][1], 15.0) self.assertLessEqual(seen[1][1], 45.0) self.assertLessEqual(seen[2][1], 180.0) def test_llm_timeout_still_returns_a_conclusion(self): fact_lines = [ "worker_nodes: titan-01, titan-02, titan-03", "hottest_cpu: longhorn-system (6.69)", ] with ( mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines), mock.patch.object(self.bot, "_ollama_call", side_effect=TimeoutError("simulated timeout")), ): reply = self.bot.open_ended_answer( "what stands out about the worker nodes?", inventory=[], snapshot=None, workloads=[], history_lines=[], mode="genius", allow_tools=True, ) self.assertIn("worker nodes", reply.lower()) self.assertIn("Confidence:", reply) def test_longhorn_rpi4_subset_beats_generic_rpi4_list(self): fact_lines = [ "rpi4: titan-12, titan-13, titan-14, titan-15, titan-17, titan-18, titan-19", "rpi4 armbian longhorn: titan-13/15/17/19", ] with ( mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines), mock.patch.object(self.bot, "_ollama_call", side_effect=AssertionError("LLM should not be called")), ): reply = self.bot.open_ended_answer( "which nodes in titan are the rpi4 longhorn nodes?", inventory=[], snapshot=None, workloads=[], history_lines=[], mode="smart", allow_tools=True, ) self.assertIn("titan-13", reply) self.assertIn("titan-15", reply) self.assertIn("titan-17", reply) self.assertIn("titan-19", reply) self.assertNotIn("titan-12", reply) self.assertNotIn("titan-14", reply) self.assertNotIn("titan-18", reply) self.assertIn("Confidence:", reply) def test_longhorn_query_uses_kb_context_when_snapshot_only_has_generic_rpi4(self): fact_lines = [ "rpi4: titan-12, titan-13, titan-14, titan-15, titan-17, titan-18, titan-19", ] kb_detail = ( "Atlas KB (retrieved):\n" "- metis (knowledge/metis.md)\n" "rpi4 armbian longhorn: titan-13/15/17/19" ) with ( mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines), mock.patch.object(self.bot, "kb_retrieve", return_value=kb_detail), mock.patch.object(self.bot, "_ollama_call", side_effect=AssertionError("LLM should not be called")), ): reply = self.bot.open_ended_answer( "which nodes in titan are the rpi4 longhorn nodes?", inventory=[], snapshot=None, workloads=[], history_lines=[], mode="smart", allow_tools=True, ) self.assertIn("titan-13", reply) self.assertIn("titan-15", reply) self.assertIn("titan-17", reply) self.assertIn("titan-19", reply) self.assertNotIn("titan-12", reply) self.assertNotIn("titan-14", reply) self.assertNotIn("titan-18", reply) self.assertIn("Confidence:", reply)