import unittest
from unittest.mock import patch, MagicMock
from universeos.client import UniverseClient, init_universe, get_client

class TestUniverseClient(unittest.TestCase):
    def setUp(self):
        self.gateway_url = "http://test-gateway:8080"
        self.client = UniverseClient(gateway_url=self.gateway_url)

    @patch('universeos.client.requests.post')
    def test_chat_completion_success(self, mock_post):
        # Setup mock
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"id": "test-id", "choices": []}
        mock_post.return_value = mock_response

        # Execute
        messages = [{"role": "user", "content": "hello"}]
        response = self.client.chat_completion(model="gpt-4", messages=messages)

        # Verify
        self.assertEqual(response["id"], "test-id")
        mock_post.assert_called_once()
        args, kwargs = mock_post.call_args
        self.assertEqual(args[0], f"{self.gateway_url}/v1/chat/completions")
        self.assertEqual(kwargs['json']['model'], "gpt-4")
        self.assertEqual(kwargs['json']['messages'], messages)

    @patch('universeos.client.requests.post')
    def test_chat_completion_failure(self, mock_post):
        # Setup mock to raise exception
        mock_post.side_effect = Exception("Connection Error")

        # Execute & Verify
        with self.assertRaises(Exception):
            self.client.chat_completion(model="gpt-4", messages=[])

    def test_init_universe(self):
        init_universe(gateway_url="http://custom-url")
        client = get_client()
        self.assertEqual(client.gateway_url, "http://custom-url")

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