62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
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")
|