import unittest import urllib.error from unittest.mock import MagicMock, patch 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()