import sqlite3 import tempfile import unittest from contextlib import closing from datetime import datetime, timezone from pathlib import Path from zoneinfo import ZoneInfo from assistant_bot.agent import execute_agent_tool, parse_agent_decision from assistant_bot.reminders import parse_reminder, unit_to_timedelta from assistant_bot.storage import AssistantStorage class ReminderParserTests(unittest.TestCase): def test_supported_relative_unit(self) -> None: delta = unit_to_timedelta(2, "часа") self.assertIsNotNone(delta) assert delta is not None self.assertEqual(delta.total_seconds(), 7200) def test_absolute_datetime(self) -> None: result = parse_reminder( "2030-01-02 10:30 проверить отчет", ZoneInfo("UTC"), ) self.assertIsNotNone(result) assert result is not None self.assertEqual( result.remind_at_utc, datetime(2030, 1, 2, 10, 30, tzinfo=timezone.utc), ) self.assertEqual(result.text, "проверить отчет") class AgentDecisionTests(unittest.TestCase): def test_tool_call_json(self) -> None: decision = parse_agent_decision( '{"tool_calls":[{"name":"create_note","arguments":{"text":"идея"}}]}' ) self.assertIsNone(decision.final) self.assertEqual(decision.tool_calls[0]["name"], "create_note") self.assertEqual(decision.tool_calls[0]["arguments"]["text"], "идея") def test_context_reset_flag(self) -> None: decision = parse_agent_decision( '{"final":"Перейдем к новой теме","reset_context":true}' ) self.assertEqual(decision.final, "Перейдем к новой теме") self.assertTrue(decision.reset_context) class StorageTests(unittest.TestCase): def test_legacy_user_settings_gets_yandex_model_column(self) -> None: with tempfile.TemporaryDirectory() as directory: database_path = Path(directory) / "assistant.sqlite3" with closing(sqlite3.connect(database_path)) as connection: connection.execute( """ CREATE TABLE user_settings ( user_id INTEGER PRIMARY KEY, ollama_model TEXT, updated_at TEXT NOT NULL ) """ ) connection.commit() storage = AssistantStorage(database_path) storage.set_user_model(42, "yandexgpt", "yandex") self.assertEqual(storage.get_user_model(42, "yandex"), "yandexgpt") def test_provider_models_are_stored_separately(self) -> None: with tempfile.TemporaryDirectory() as directory: storage = AssistantStorage(Path(directory) / "assistant.sqlite3") storage.set_user_model(42, "qwen3.5:9b", "local") storage.set_user_model(42, "yandexgpt", "yandex") self.assertEqual(storage.get_user_model(42, "local"), "qwen3.5:9b") self.assertEqual(storage.get_user_model(42, "yandex"), "yandexgpt") def test_memory_crud(self) -> None: with tempfile.TemporaryDirectory() as directory: storage = AssistantStorage(Path(directory) / "assistant.sqlite3") memory_id = storage.add_memory(42, "короткие ответы") self.assertEqual(storage.list_memories(42)[0]["id"], memory_id) self.assertTrue(storage.delete_memory(42, memory_id)) self.assertEqual(storage.list_memories(42), []) def test_new_context_preserves_searchable_archive(self) -> None: with tempfile.TemporaryDirectory() as directory: storage = AssistantStorage(Path(directory) / "assistant.sqlite3") storage.add_conversation_exchange( 42, 100, "Обсудим проект Альфа", "Какой аспект проекта интересует?", ) storage.start_new_conversation(42, 100) storage.add_conversation_exchange( 42, 100, "Снова обсуждаем проект Альфа", "Продолжаем обсуждение проекта.", ) storage.add_conversation_exchange(42, 200, "Другой чат", "Другой ответ") messages = storage.list_conversation_messages(42, 100) matches = storage.search_conversation_messages(42, 100, "проект Альфа") self.assertEqual( [(row["role"], row["content"]) for row in messages], [ ("user", "Снова обсуждаем проект Альфа"), ("assistant", "Продолжаем обсуждение проекта."), ], ) self.assertEqual(matches[0]["content"], "Снова обсуждаем проект Альфа") self.assertIn("Обсудим проект Альфа", [row["content"] for row in matches]) self.assertEqual(len(storage.list_conversation_messages(42, 200)), 2) archive_window = storage.conversation_message_window( 42, 100, int(matches[-1]["id"]), ) self.assertIn( "Какой аспект проекта интересует?", [row["content"] for row in archive_window], ) tool_result = execute_agent_tool( "search_conversation", {"query": "проект Альфа", "limit": 2}, storage, user_id=42, chat_id=100, tz=ZoneInfo("UTC"), ) self.assertTrue(tool_result["ok"]) self.assertTrue(tool_result["discussions"]) self.assertTrue(tool_result["discussions"][0]["messages"]) if __name__ == "__main__": unittest.main()