init

tests/test_yandex_ai.py

@@ -0,0 +1,132 @@
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from assistant_bot.yandex_ai import YandexAIClient, YandexSpeechRecognizer
+
+
+class FakeCompletion:
+    def __init__(self) -> None:
+        self.configuration = None
+        self.messages = None
+        self.timeout = None
+
+    def configure(self, **kwargs):
+        self.configuration = kwargs
+        return self
+
+    async def run(self, messages, timeout):
+        self.messages = messages
+        self.timeout = timeout
+        return [SimpleNamespace(text="  готово  ")]
+
+
+class YandexAIClientTests(unittest.IsolatedAsyncioTestCase):
+    async def test_chat_uses_sdk_messages_and_json_mode(self) -> None:
+        completion = FakeCompletion()
+        requested_models = []
+
+        def completions(model):
+            requested_models.append(model)
+            return completion
+
+        sdk = SimpleNamespace(
+            models=SimpleNamespace(completions=completions),
+            chat=SimpleNamespace(
+                completions=SimpleNamespace(
+                    list=self._list_models,
+                )
+            ),
+        )
+        client = YandexAIClient(folder_id="folder-id", sdk=sdk)
+
+        answer = await client.chat(
+            "yandexgpt",
+            [
+                {"role": "system", "content": "Отвечай кратко"},
+                {"role": "user", "content": "Привет"},
+            ],
+            json_mode=True,
+        )
+
+        self.assertEqual(answer, "готово")
+        self.assertEqual(requested_models, ["gpt://folder-id/yandexgpt"])
+        self.assertEqual(completion.configuration, {"response_format": "json"})
+        self.assertEqual(
+            completion.messages,
+            [
+                {"role": "system", "text": "Отвечай кратко"},
+                {"role": "user", "text": "Привет"},
+            ],
+        )
+        self.assertEqual(completion.timeout, 180)
+
+    async def test_list_models_returns_sorted_uris(self) -> None:
+        sdk = SimpleNamespace(
+            chat=SimpleNamespace(
+                completions=SimpleNamespace(
+                    list=self._list_models,
+                )
+            )
+        )
+        client = YandexAIClient(folder_id="folder-id", sdk=sdk)
+
+        self.assertEqual(
+            await client.list_models(),
+            ["gpt://folder/alice-ai/latest", "gpt://folder/yandexgpt/latest"],
+        )
+
+    @staticmethod
+    async def _list_models():
+        return [
+            SimpleNamespace(uri="gpt://folder/yandexgpt/latest"),
+            SimpleNamespace(uri="gpt://folder/alice-ai/latest"),
+        ]
+
+
+class YandexSpeechRecognizerTests(unittest.TestCase):
+    def test_transcribe_uses_speechkit_ogg_opus(self) -> None:
+        calls = []
+
+        class FakeRecognizer:
+            def run(self, audio, timeout):
+                calls.append((audio, timeout))
+                return SimpleNamespace(text="  Привет, мир  ")
+
+        format_value = object()
+
+        def speech_to_text(**kwargs):
+            calls.append(kwargs)
+            return FakeRecognizer()
+
+        sdk = SimpleNamespace(
+            speechkit=SimpleNamespace(
+                AudioFormat=SimpleNamespace(OGG_OPUS=format_value),
+                speech_to_text=speech_to_text,
+            )
+        )
+        recognizer = YandexSpeechRecognizer(
+            folder_id="folder-id",
+            language="ru-RU",
+            model="general",
+            sdk=sdk,
+        )
+
+        with patch.object(Path, "read_bytes", return_value=b"ogg-data"):
+            result = recognizer.transcribe("voice.ogg")
+
+        self.assertEqual(result, "Привет, мир")
+        self.assertEqual(
+            calls[0],
+            {
+                "audio_format": format_value,
+                "language_codes": "ru-RU",
+                "model": "general",
+            },
+        )
+        self.assertEqual(calls[1], (b"ogg-data", 180))
+
+
+if __name__ == "__main__":
+    unittest.main()