bstein-dev-home/backend/tests/test_ai.py

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