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