|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
import tempfile
|
|
import os
|
|
from pr.core.assistant import Assistant, process_message
|
|
|
|
|
|
class TestAssistant(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.args = MagicMock()
|
|
self.args.verbose = False
|
|
self.args.debug = False
|
|
self.args.no_syntax = False
|
|
self.args.model = 'test-model'
|
|
self.args.api_url = 'test-url'
|
|
self.args.model_list_url = 'test-list-url'
|
|
|
|
@patch('sqlite3.connect')
|
|
@patch('os.environ.get')
|
|
@patch('pr.core.context.init_system_message')
|
|
@patch('pr.core.enhanced_assistant.EnhancedAssistant')
|
|
def test_init(self, mock_enhanced, mock_init_sys, mock_env, mock_sqlite):
|
|
mock_env.side_effect = lambda key, default: {'OPENROUTER_API_KEY': 'key', 'AI_MODEL': 'model', 'API_URL': 'url', 'MODEL_LIST_URL': 'list', 'USE_TOOLS': '1', 'STRICT_MODE': '0'}.get(key, default)
|
|
mock_conn = MagicMock()
|
|
mock_sqlite.return_value = mock_conn
|
|
mock_init_sys.return_value = {'role': 'system', 'content': 'sys'}
|
|
|
|
assistant = Assistant(self.args)
|
|
|
|
self.assertEqual(assistant.api_key, 'key')
|
|
self.assertEqual(assistant.model, 'test-model')
|
|
mock_sqlite.assert_called_once()
|
|
|
|
@patch('pr.core.assistant.call_api')
|
|
@patch('pr.core.assistant.render_markdown')
|
|
def test_process_response_no_tools(self, mock_render, mock_call):
|
|
assistant = MagicMock()
|
|
assistant.messages = MagicMock()
|
|
assistant.verbose = False
|
|
assistant.syntax_highlighting = True
|
|
mock_render.return_value = 'rendered'
|
|
|
|
response = {'choices': [{'message': {'content': 'content'}}]}
|
|
|
|
result = Assistant.process_response(assistant, response)
|
|
|
|
self.assertEqual(result, 'rendered')
|
|
assistant.messages.append.assert_called_with({'content': 'content'})
|
|
|
|
@patch('pr.core.assistant.call_api')
|
|
@patch('pr.core.assistant.render_markdown')
|
|
@patch('pr.core.assistant.get_tools_definition')
|
|
def test_process_response_with_tools(self, mock_tools_def, mock_render, mock_call):
|
|
assistant = MagicMock()
|
|
assistant.messages = MagicMock()
|
|
assistant.verbose = False
|
|
assistant.syntax_highlighting = True
|
|
assistant.use_tools = True
|
|
assistant.model = 'model'
|
|
assistant.api_url = 'url'
|
|
assistant.api_key = 'key'
|
|
mock_tools_def.return_value = []
|
|
mock_call.return_value = {'choices': [{'message': {'content': 'follow'}}]}
|
|
|
|
response = {'choices': [{'message': {'tool_calls': [{'id': '1', 'function': {'name': 'test', 'arguments': '{}'}}]}}]}
|
|
|
|
with patch.object(assistant, 'execute_tool_calls', return_value=[{'role': 'tool', 'content': 'result'}]):
|
|
result = Assistant.process_response(assistant, response)
|
|
|
|
mock_call.assert_called()
|
|
|
|
@patch('pr.core.assistant.call_api')
|
|
@patch('pr.core.assistant.get_tools_definition')
|
|
def test_process_message(self, mock_tools, mock_call):
|
|
assistant = MagicMock()
|
|
assistant.messages = MagicMock()
|
|
assistant.verbose = False
|
|
assistant.use_tools = True
|
|
assistant.model = 'model'
|
|
assistant.api_url = 'url'
|
|
assistant.api_key = 'key'
|
|
mock_tools.return_value = []
|
|
mock_call.return_value = {'choices': [{'message': {'content': 'response'}}]}
|
|
|
|
with patch('pr.core.assistant.render_markdown', return_value='rendered'):
|
|
with patch('builtins.print'):
|
|
process_message(assistant, 'test message')
|
|
|
|
assistant.messages.append.assert_called_with({'role': 'user', 'content': 'test message'})
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|