import unittest
from unittest.mock import patch

from assistant_bot.ollama import OllamaClient


class OllamaClientTests(unittest.TestCase):
    def test_chat_logs_token_and_timing_metrics(self) -> None:
        client = OllamaClient("http://localhost:11434")
        response = {
            "message": {"content": "  готово  "},
            "prompt_eval_count": 2004,
            "eval_count": 12,
            "prompt_eval_duration": 91_400_000,
            "eval_duration": 25_000_000,
            "total_duration": 130_000_000,
            "load_duration": 4_000_000,
        }

        with (
            patch.object(
                client,
                "_request_json",
                return_value=response,
            ) as request_json,
            self.assertLogs(
                "assistant_bot.ollama",
                level="INFO",
            ) as captured,
        ):
            answer = client._chat(
                "qwen3.5:9b",
                [{"role": "user", "content": "Привет"}],
                json_mode=True,
            )

        self.assertEqual(answer, "готово")
        payload = request_json.call_args.args[2]
        self.assertEqual(payload["format"], "json")
        self.assertIn("input_tokens=2004", captured.output[0])
        self.assertIn("output_tokens=12", captured.output[0])
        self.assertIn("total_tokens=2016", captured.output[0])
        self.assertIn("prompt_eval_ms=91.4", captured.output[0])


if __name__ == "__main__":
    unittest.main()