kandrusyak_bot/tests/test_yandex_ai.py

import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

from assistant_bot.yandex_ai import (
    YandexAIClient,
    YandexSpeechRecognizer,
    describe_yandex_error,
    prepare_chat_messages,
)


class FakeCompletion:
    def __init__(self) -> None:
        self.configuration = None
        self.messages = None
        self.timeout = None
        self.result = [SimpleNamespace(text="  готово  ")]

    def configure(self, **kwargs):
        self.configuration = kwargs
        return self

    async def run(self, messages, timeout):
        self.messages = messages
        self.timeout = timeout
        return self.result


class FakeChatCompletions:
    def __init__(self, completion, list_models) -> None:
        self.completion = completion
        self.list = list_models
        self.requested_models = []

    def __call__(self, model):
        self.requested_models.append(model)
        return self.completion


class FakeHTTPError(RuntimeError):
    def __init__(self) -> None:
        super().__init__("forbidden")
        self.response = SimpleNamespace(
            status_code=403,
            text='{"error":{"message":"folder mismatch"}}',
        )


class YandexAIClientTests(unittest.IsolatedAsyncioTestCase):
    def test_multiple_system_messages_are_merged_at_the_beginning(self) -> None:
        self.assertEqual(
            prepare_chat_messages(
                [
                    {"role": "system", "content": "Primary instructions"},
                    {"role": "user", "content": "Earlier question"},
                    {"role": "system", "content": "Runtime context"},
                    {"role": "assistant", "content": "Earlier answer"},
                    {"role": "user", "content": "Current question"},
                ]
            ),
            [
                {
                    "role": "system",
                    "content": "Primary instructions\n\nRuntime context",
                },
                {"role": "user", "content": "Earlier question"},
                {"role": "assistant", "content": "Earlier answer"},
                {"role": "user", "content": "Current question"},
            ],
        )

    def test_http_error_includes_safe_response_body(self) -> None:
        error = FakeHTTPError()

        self.assertEqual(
            describe_yandex_error(error),
            'HTTP 403: {"error":{"message":"folder mismatch"}}',
        )

    async def test_chat_uses_sdk_messages_and_json_mode(self) -> None:
        completion = FakeCompletion()
        completions = FakeChatCompletions(completion, self._list_models)

        sdk = SimpleNamespace(
            chat=SimpleNamespace(
                completions=completions,
            ),
        )
        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(
            completions.requested_models,
            ["gpt://folder-id/yandexgpt"],
        )
        self.assertEqual(completion.configuration, {"response_format": "json"})
        self.assertEqual(
            completion.messages,
            [
                {"role": "system", "content": "Отвечай кратко"},
                {"role": "user", "content": "Привет"},
            ],
        )
        self.assertEqual(completion.timeout, 180)

    async def test_chat_repairs_legacy_uri_without_folder_id(self) -> None:
        completion = FakeCompletion()
        completions = FakeChatCompletions(completion, self._list_models)
        sdk = SimpleNamespace(
            chat=SimpleNamespace(completions=completions),
        )
        client = YandexAIClient(folder_id="folder-id", sdk=sdk)

        await client.chat(
            "gpt://qwen3.6-35b-a3b/latest",
            [{"role": "user", "content": "Привет"}],
        )

        self.assertEqual(
            completions.requested_models,
            ["gpt://folder-id/qwen3.6-35b-a3b/latest"],
        )

    async def test_chat_logs_reported_token_usage(self) -> None:
        class FakeResult(list):
            pass

        completion = FakeCompletion()
        completion.result = FakeResult(
            [SimpleNamespace(text="готово")]
        )
        completion.result.usage = SimpleNamespace(
            prompt_tokens=120,
            completion_tokens=8,
            total_tokens=128,
        )
        completions = FakeChatCompletions(completion, self._list_models)
        sdk = SimpleNamespace(
            chat=SimpleNamespace(completions=completions),
        )
        client = YandexAIClient(folder_id="folder-id", sdk=sdk)

        with self.assertLogs(
            "assistant_bot.yandex_ai",
            level="INFO",
        ) as captured:
            await client.chat(
                "yandexgpt",
                [{"role": "user", "content": "Привет"}],
            )

        self.assertIn("input_tokens=120", captured.output[0])
        self.assertIn("output_tokens=8", captured.output[0])
        self.assertIn("total_tokens=128", captured.output[0])

    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()