kandrusyak_bot/tests/test_core.py

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_user_authorization_is_persisted(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            database_path = Path(directory) / "assistant.sqlite3"
            storage = AssistantStorage(database_path)

            self.assertFalse(storage.is_user_authorized(42))
            storage.authorize_user(42)

            reopened_storage = AssistantStorage(database_path)
            self.assertTrue(reopened_storage.is_user_authorized(42))
            self.assertFalse(reopened_storage.is_user_authorized(43))

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