chore: remove verbose prints and add agent/memory tool registration in assistant core
This commit is contained in:
@@ -0,0 +1,81 @@
|
||||
import pytest
|
||||
from pr.core.advanced_context import AdvancedContextManager
|
||||
|
||||
def test_adaptive_context_window_simple():
|
||||
mgr = AdvancedContextManager()
|
||||
messages = [{'content': 'short'}, {'content': 'this is a longer message with more words'}]
|
||||
window = mgr.adaptive_context_window(messages, 'simple')
|
||||
assert isinstance(window, int)
|
||||
assert window >= 10
|
||||
|
||||
def test_adaptive_context_window_medium():
|
||||
mgr = AdvancedContextManager()
|
||||
messages = [{'content': 'short'}, {'content': 'this is a longer message with more words'}]
|
||||
window = mgr.adaptive_context_window(messages, 'medium')
|
||||
assert isinstance(window, int)
|
||||
assert window >= 20
|
||||
|
||||
def test_adaptive_context_window_complex():
|
||||
mgr = AdvancedContextManager()
|
||||
messages = [{'content': 'short'}, {'content': 'this is a longer message with more words'}]
|
||||
window = mgr.adaptive_context_window(messages, 'complex')
|
||||
assert isinstance(window, int)
|
||||
assert window >= 35
|
||||
|
||||
def test_analyze_message_complexity():
|
||||
mgr = AdvancedContextManager()
|
||||
messages = [{'content': 'hello world'}, {'content': 'hello again'}]
|
||||
score = mgr._analyze_message_complexity(messages)
|
||||
assert 0 <= score <= 1
|
||||
|
||||
def test_analyze_message_complexity_empty():
|
||||
mgr = AdvancedContextManager()
|
||||
messages = []
|
||||
score = mgr._analyze_message_complexity(messages)
|
||||
assert score == 0
|
||||
|
||||
def test_extract_key_sentences():
|
||||
mgr = AdvancedContextManager()
|
||||
text = "This is the first sentence. This is the second sentence. This is a longer third sentence with more words."
|
||||
sentences = mgr.extract_key_sentences(text, 2)
|
||||
assert len(sentences) <= 2
|
||||
assert all(isinstance(s, str) for s in sentences)
|
||||
|
||||
def test_extract_key_sentences_empty():
|
||||
mgr = AdvancedContextManager()
|
||||
text = ""
|
||||
sentences = mgr.extract_key_sentences(text, 5)
|
||||
assert sentences == []
|
||||
|
||||
def test_advanced_summarize_messages():
|
||||
mgr = AdvancedContextManager()
|
||||
messages = [{'content': 'Hello'}, {'content': 'How are you?'}]
|
||||
summary = mgr.advanced_summarize_messages(messages)
|
||||
assert isinstance(summary, str)
|
||||
|
||||
def test_advanced_summarize_messages_empty():
|
||||
mgr = AdvancedContextManager()
|
||||
messages = []
|
||||
summary = mgr.advanced_summarize_messages(messages)
|
||||
assert summary == "No content to summarize."
|
||||
|
||||
def test_score_message_relevance():
|
||||
mgr = AdvancedContextManager()
|
||||
message = {'content': 'hello world'}
|
||||
context = 'world hello'
|
||||
score = mgr.score_message_relevance(message, context)
|
||||
assert 0 <= score <= 1
|
||||
|
||||
def test_score_message_relevance_no_overlap():
|
||||
mgr = AdvancedContextManager()
|
||||
message = {'content': 'hello'}
|
||||
context = 'world'
|
||||
score = mgr.score_message_relevance(message, context)
|
||||
assert score == 0
|
||||
|
||||
def test_score_message_relevance_empty():
|
||||
mgr = AdvancedContextManager()
|
||||
message = {'content': ''}
|
||||
context = ''
|
||||
score = mgr.score_message_relevance(message, context)
|
||||
assert score == 0
|
||||
@@ -0,0 +1,63 @@
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import json
|
||||
import urllib.error
|
||||
from pr.core.api import call_api, list_models
|
||||
|
||||
|
||||
class TestApi(unittest.TestCase):
|
||||
|
||||
@patch('pr.core.api.urllib.request.urlopen')
|
||||
@patch('pr.core.api.auto_slim_messages')
|
||||
def test_call_api_success(self, mock_slim, mock_urlopen):
|
||||
mock_slim.return_value = [{'role': 'user', 'content': 'test'}]
|
||||
mock_response = MagicMock()
|
||||
mock_response.read.return_value = b'{"choices": [{"message": {"content": "response"}}], "usage": {"tokens": 10}}'
|
||||
mock_urlopen.return_value.__enter__.return_value = mock_response
|
||||
|
||||
result = call_api([], 'model', 'http://url', 'key', True, [{'name': 'tool'}])
|
||||
|
||||
self.assertIn('choices', result)
|
||||
mock_urlopen.assert_called_once()
|
||||
|
||||
@patch('urllib.request.urlopen')
|
||||
@patch('pr.core.api.auto_slim_messages')
|
||||
def test_call_api_http_error(self, mock_slim, mock_urlopen):
|
||||
mock_slim.return_value = [{'role': 'user', 'content': 'test'}]
|
||||
mock_urlopen.side_effect = urllib.error.HTTPError('http://url', 500, 'error', None, MagicMock())
|
||||
|
||||
result = call_api([], 'model', 'http://url', 'key', False, [])
|
||||
|
||||
self.assertIn('error', result)
|
||||
|
||||
@patch('urllib.request.urlopen')
|
||||
@patch('pr.core.api.auto_slim_messages')
|
||||
def test_call_api_general_error(self, mock_slim, mock_urlopen):
|
||||
mock_slim.return_value = [{'role': 'user', 'content': 'test'}]
|
||||
mock_urlopen.side_effect = Exception('test error')
|
||||
|
||||
result = call_api([], 'model', 'http://url', 'key', False, [])
|
||||
|
||||
self.assertIn('error', result)
|
||||
|
||||
@patch('urllib.request.urlopen')
|
||||
def test_list_models_success(self, mock_urlopen):
|
||||
mock_response = MagicMock()
|
||||
mock_response.read.return_value = b'{"data": [{"id": "model1"}]}'
|
||||
mock_urlopen.return_value.__enter__.return_value = mock_response
|
||||
|
||||
result = list_models('http://url', 'key')
|
||||
|
||||
self.assertEqual(result, [{'id': 'model1'}])
|
||||
|
||||
@patch('urllib.request.urlopen')
|
||||
def test_list_models_error(self, mock_urlopen):
|
||||
mock_urlopen.side_effect = Exception('error')
|
||||
|
||||
result = list_models('http://url', 'key')
|
||||
|
||||
self.assertIn('error', result)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,94 @@
|
||||
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()
|
||||
@@ -0,0 +1,56 @@
|
||||
import pytest
|
||||
from unittest.mock import patch, mock_open
|
||||
import os
|
||||
from pr.core.config_loader import load_config, _load_config_file, _parse_value, create_default_config
|
||||
|
||||
def test_parse_value_string():
|
||||
assert _parse_value('hello') == 'hello'
|
||||
|
||||
def test_parse_value_int():
|
||||
assert _parse_value('123') == 123
|
||||
|
||||
def test_parse_value_float():
|
||||
assert _parse_value('1.23') == 1.23
|
||||
|
||||
def test_parse_value_bool_true():
|
||||
assert _parse_value('true') == True
|
||||
|
||||
def test_parse_value_bool_false():
|
||||
assert _parse_value('false') == False
|
||||
|
||||
def test_parse_value_bool_upper():
|
||||
assert _parse_value('TRUE') == True
|
||||
|
||||
@patch('os.path.exists', return_value=False)
|
||||
def test_load_config_file_not_exists(mock_exists):
|
||||
config = _load_config_file('test.ini')
|
||||
assert config == {}
|
||||
|
||||
@patch('os.path.exists', return_value=True)
|
||||
@patch('configparser.ConfigParser')
|
||||
def test_load_config_file_exists(mock_parser_class, mock_exists):
|
||||
mock_parser = mock_parser_class.return_value
|
||||
mock_parser.sections.return_value = ['api']
|
||||
mock_parser.items.return_value = [('key', 'value')]
|
||||
config = _load_config_file('test.ini')
|
||||
assert 'api' in config
|
||||
assert config['api']['key'] == 'value'
|
||||
|
||||
@patch('pr.core.config_loader._load_config_file')
|
||||
def test_load_config(mock_load):
|
||||
mock_load.side_effect = [{'api': {'key': 'global'}}, {'api': {'key': 'local'}}]
|
||||
config = load_config()
|
||||
assert config['api']['key'] == 'local'
|
||||
|
||||
@patch('builtins.open', new_callable=mock_open)
|
||||
def test_create_default_config(mock_file):
|
||||
result = create_default_config('test.ini')
|
||||
assert result == True
|
||||
mock_file.assert_called_once_with('test.ini', 'w')
|
||||
handle = mock_file()
|
||||
handle.write.assert_called_once()
|
||||
|
||||
@patch('builtins.open', side_effect=Exception('error'))
|
||||
def test_create_default_config_error(mock_file):
|
||||
result = create_default_config('test.ini')
|
||||
assert result == False
|
||||
@@ -0,0 +1,118 @@
|
||||
import pytest
|
||||
from unittest.mock import patch
|
||||
import sys
|
||||
from pr.__main__ import main
|
||||
|
||||
def test_main_version(capsys):
|
||||
with patch('sys.argv', ['pr', '--version']):
|
||||
with pytest.raises(SystemExit):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'PR Assistant' in captured.out
|
||||
|
||||
def test_main_create_config_success(capsys):
|
||||
with patch('pr.core.config_loader.create_default_config', return_value=True):
|
||||
with patch('sys.argv', ['pr', '--create-config']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'Configuration file created' in captured.out
|
||||
|
||||
def test_main_create_config_fail(capsys):
|
||||
with patch('pr.core.config_loader.create_default_config', return_value=False):
|
||||
with patch('sys.argv', ['pr', '--create-config']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'Error creating configuration file' in captured.err
|
||||
|
||||
def test_main_list_sessions_no_sessions(capsys):
|
||||
with patch('pr.core.session.SessionManager') as mock_sm:
|
||||
mock_instance = mock_sm.return_value
|
||||
mock_instance.list_sessions.return_value = []
|
||||
with patch('sys.argv', ['pr', '--list-sessions']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'No saved sessions found' in captured.out
|
||||
|
||||
def test_main_list_sessions_with_sessions(capsys):
|
||||
sessions = [{'name': 'test', 'created_at': '2023-01-01', 'message_count': 5}]
|
||||
with patch('pr.core.session.SessionManager') as mock_sm:
|
||||
mock_instance = mock_sm.return_value
|
||||
mock_instance.list_sessions.return_value = sessions
|
||||
with patch('sys.argv', ['pr', '--list-sessions']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'Found 1 saved sessions' in captured.out
|
||||
assert 'test' in captured.out
|
||||
|
||||
def test_main_delete_session_success(capsys):
|
||||
with patch('pr.core.session.SessionManager') as mock_sm:
|
||||
mock_instance = mock_sm.return_value
|
||||
mock_instance.delete_session.return_value = True
|
||||
with patch('sys.argv', ['pr', '--delete-session', 'test']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert "Session 'test' deleted" in captured.out
|
||||
|
||||
def test_main_delete_session_fail(capsys):
|
||||
with patch('pr.core.session.SessionManager') as mock_sm:
|
||||
mock_instance = mock_sm.return_value
|
||||
mock_instance.delete_session.return_value = False
|
||||
with patch('sys.argv', ['pr', '--delete-session', 'test']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert "Error deleting session 'test'" in captured.err
|
||||
|
||||
def test_main_export_session_json(capsys):
|
||||
with patch('pr.core.session.SessionManager') as mock_sm:
|
||||
mock_instance = mock_sm.return_value
|
||||
mock_instance.export_session.return_value = True
|
||||
with patch('sys.argv', ['pr', '--export-session', 'test', 'output.json']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'Session exported to output.json' in captured.out
|
||||
|
||||
def test_main_export_session_md(capsys):
|
||||
with patch('pr.core.session.SessionManager') as mock_sm:
|
||||
mock_instance = mock_sm.return_value
|
||||
mock_instance.export_session.return_value = True
|
||||
with patch('sys.argv', ['pr', '--export-session', 'test', 'output.md']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'Session exported to output.md' in captured.out
|
||||
|
||||
def test_main_usage(capsys):
|
||||
usage = {'total_requests': 10, 'total_tokens': 1000, 'total_cost': 0.01}
|
||||
with patch('pr.core.usage_tracker.UsageTracker.get_total_usage', return_value=usage):
|
||||
with patch('sys.argv', ['pr', '--usage']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'Total Usage Statistics' in captured.out
|
||||
assert 'Requests: 10' in captured.out
|
||||
|
||||
def test_main_plugins_no_plugins(capsys):
|
||||
with patch('pr.plugins.loader.PluginLoader') as mock_loader:
|
||||
mock_instance = mock_loader.return_value
|
||||
mock_instance.load_plugins.return_value = None
|
||||
mock_instance.list_loaded_plugins.return_value = []
|
||||
with patch('sys.argv', ['pr', '--plugins']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'No plugins loaded' in captured.out
|
||||
|
||||
def test_main_plugins_with_plugins(capsys):
|
||||
with patch('pr.plugins.loader.PluginLoader') as mock_loader:
|
||||
mock_instance = mock_loader.return_value
|
||||
mock_instance.load_plugins.return_value = None
|
||||
mock_instance.list_loaded_plugins.return_value = ['plugin1', 'plugin2']
|
||||
with patch('sys.argv', ['pr', '--plugins']):
|
||||
main()
|
||||
captured = capsys.readouterr()
|
||||
assert 'Loaded 2 plugins' in captured.out
|
||||
|
||||
def test_main_run_assistant():
|
||||
with patch('pr.__main__.Assistant') as mock_assistant:
|
||||
mock_instance = mock_assistant.return_value
|
||||
with patch('sys.argv', ['pr', 'test message']):
|
||||
main()
|
||||
mock_assistant.assert_called_once()
|
||||
mock_instance.run.assert_called_once()
|
||||
Reference in New Issue
Block a user