tests/test_gigachat_assistant.py - main - cppv/hackathon13

tests/test_gigachat_assistant.py
115 строк4 KB
cppv
hackathon13
25 май 2025, 11:54
25 май 2025, 11:548af8c2d
import unittest
from unittest.mock import patch, MagicMock
import sys
import os

# чтоб работало
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from app.gigachat_assistant import GigaChatAssistant

class TestGigaChatAssistant(unittest.TestCase):
    """
    тесты для гигачата
    """
    
    def setUp(self):
        """
        подготовка всякого
        """
        self.credentials = {"auth_key": "test_key"}
        self.mock_vectorstore = MagicMock()
        self.mock_vectorstore.as_retriever.return_value = MagicMock()
        
    @patch('app.gigachat_assistant.GigaChat')
    def test_initialization(self, mock_gigachat):
        """
        проверяем что все запускается
        """
        # делаем мок
        mock_gigachat.return_value = MagicMock()
        
        # создаем штуку
        assistant = GigaChatAssistant(self.credentials, self.mock_vectorstore)
        
        # проверяем что вызвалось норм
        mock_gigachat.assert_called_once_with(
            credentials=self.credentials["auth_key"],  # тут надо строку а не словарь
            verify_ssl_certs=False
        )
        
        # и ретривер создался
        self.mock_vectorstore.as_retriever.assert_called_once()
        
    @patch('app.gigachat_assistant.GigaChat')
    def test_answer_question(self, mock_gigachat):
        """
        проверяем ответы на вопросы
        """
        # моки всякие
        mock_chat_model = MagicMock()
        mock_gigachat.return_value = mock_chat_model
        
        # еще моки
        mock_qa_chain = MagicMock()
        mock_qa_chain.run.return_value = "Тестовый ответ"
        
        # делаем ассистента
        assistant = GigaChatAssistant(self.credentials, self.mock_vectorstore)
        assistant.qa_chain = mock_qa_chain
        
        # спрашиваем чет
        response = assistant.answer_question("Тестовый вопрос")
        
        # проверяем что вызвалось
        mock_qa_chain.run.assert_called_once_with("Тестовый вопрос")
        
        # и ответ норм
        self.assertEqual(response, "Тестовый ответ")
        
    @patch('app.gigachat_assistant.GigaChat')
    def test_direct_chat(self, mock_gigachat):
        """
        тест прямого чата
        """
        # моки опять
        mock_chat_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "Тестовый ответ"
        mock_chat_model.return_value = mock_response
        mock_gigachat.return_value = mock_chat_model
        
        # создаем
        assistant = GigaChatAssistant(self.credentials, self.mock_vectorstore)
        
        # чатимся
        response = assistant.direct_chat("Тестовый вопрос")
        
        # проверяем что работает
        mock_chat_model.assert_called_once()
        
        # и ответ норм
        self.assertEqual(response, "Тестовый ответ")
        
    @patch('app.gigachat_assistant.GigaChat')
    def test_error_handling(self, mock_gigachat):
        """
        проверяем что будет если все сломается
        """
        # делаем чтоб была ошибка
        mock_chat_model = MagicMock()
        mock_chat_model.side_effect = Exception("Тестовая ошибка")
        mock_gigachat.return_value = mock_chat_model
        
        # создаем
        assistant = GigaChatAssistant(self.credentials, self.mock_vectorstore)
        
        # пробуем чатиться
        response = assistant.direct_chat("Тестовый вопрос")
        
        # смотрим что ошибка норм показывается
        self.assertIn("Произошла ошибка", response)
        self.assertIn("Тестовая ошибка", response)

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