48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
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()
|