from __future__ import annotations import json from unittest import TestCase, mock from atlas_portal.app_factory import create_app from atlas_portal.routes import ai class AiRouteTests(TestCase): @classmethod def setUpClass(cls): cls.app = create_app() cls.client = cls.app.test_client() def test_chat_routes_profiles_to_modes(self): seen: list[tuple[str, str]] = [] def fake_atlasbot_answer(message: str, mode: str, conversation_id: str) -> str: seen.append((mode, conversation_id)) return f"{mode}:{conversation_id}" with mock.patch.object(ai, "_atlasbot_answer", side_effect=fake_atlasbot_answer): for profile, expected_mode in ( ("atlas-quick", "quick"), ("atlas-smart", "smart"), ("atlas-genius", "genius"), ): resp = self.client.post( "/api/chat", data=json.dumps( { "message": "How is Titan doing?", "profile": profile, "conversation_id": f"conv-{profile}", } ), content_type="application/json", ) data = resp.get_json() self.assertEqual(resp.status_code, 200) self.assertEqual(data.get("source"), f"atlas-{expected_mode}") self.assertEqual(data.get("reply"), f"{expected_mode}:conv-{profile}") self.assertEqual( seen, [ ("quick", "conv-atlas-quick"), ("smart", "conv-atlas-smart"), ("genius", "conv-atlas-genius"), ], ) def test_info_endpoint_exposes_profile_specific_model(self): with mock.patch.object(ai.settings, "AI_ATLASBOT_MODEL_GENIUS", "genius-model"): resp = self.client.get("/api/ai/info?profile=atlas-genius") data = resp.get_json() self.assertEqual(resp.status_code, 200) self.assertEqual(data.get("profile"), "atlas-genius") self.assertEqual(data.get("model"), "genius-model") def test_atlasbot_answer_uses_profile_specific_timeout(self): captured: dict[str, object] = {} class DummyResponse: status_code = 200 def json(self): return {"reply": "atlas reply"} class DummyClient: def __init__(self, timeout): captured["timeout"] = timeout def __enter__(self): return self def __exit__(self, exc_type, exc, tb): return False def post(self, endpoint, json=None, headers=None): captured["endpoint"] = endpoint captured["json"] = json captured["headers"] = headers or {} return DummyResponse() with ( mock.patch.object(ai.httpx, "Client", DummyClient), mock.patch.object(ai.settings, "AI_ATLASBOT_ENDPOINT", "http://atlasbot.invalid/v1/answer"), ): reply = ai._atlasbot_answer("How is Titan doing?", "genius", "conv-1") self.assertEqual(reply, "atlas reply") self.assertEqual(captured["timeout"], ai.settings.AI_ATLASBOT_TIMEOUT_GENIUS_SEC) self.assertEqual(captured["json"], {"prompt": "How is Titan doing?", "mode": "genius", "conversation_id": "conv-1"}) def test_chat_returns_fallback_when_atlasbot_returns_empty(self): with mock.patch.object(ai, "_atlasbot_answer", return_value=""): resp = self.client.post( "/api/chat", data=json.dumps( { "message": "How is Titan doing?", "profile": "atlas-quick", } ), content_type="application/json", ) data = resp.get_json() self.assertEqual(resp.status_code, 200) self.assertEqual(data.get("source"), "atlas-quick") self.assertIn("Quick mode hit", data.get("reply", ""))