import unittest
from unittest.mock import patch, MagicMock
import os
from litellm.genesis_client import GenesisClient

class GenesisClientTest(unittest.TestCase):

    @patch('litellm.completion')
    def test_chat_completion_success(self, mock_completion):
        mock_completion.return_value = MagicMock(choices=[MagicMock(message=MagicMock(content='Test response'))])
        client = GenesisClient()
        messages = [{'role': 'user', 'content': 'Test message'}]
        response = client.chat_completion(messages=messages, model='gpt-3.5-turbo')
        self.assertIsNotNone(response)

    @patch('litellm.embedding')
    def test_embedding_success(self, mock_embedding):
        mock_embedding.return_value = MagicMock(data=[MagicMock(embedding=[0.1, 0.2, 0.3])])
        client = GenesisClient()
        response = client.embedding(input='Test input', model='text-embedding-ada-002')
        self.assertIsNotNone(response)

    @patch('litellm.completion', side_effect=Exception('Test exception'))
    def test_chat_completion_failure(self, mock_completion):
        client = GenesisClient()
        messages = [{'role': 'user', 'content': 'Test message'}]
        with self.assertRaises(Exception):
            client.chat_completion(messages=messages, model='gpt-3.5-turbo')

    @patch('litellm.embedding', side_effect=Exception('Test exception'))
    def test_embedding_failure(self, mock_embedding):
        client = GenesisClient()
        with self.assertRaises(Exception):
            client.embedding(input='Test input', model='text-embedding-ada-002')

    def test_model_config(self):
        model_config = {
            'gpt-3.5-turbo': {'temperature': 0.5},
            'claude-v1.3': {'max_tokens': 100}
        }
        client = GenesisClient(model_config=model_config)
        self.assertEqual(client.model_config, model_config)

    def test_default_model(self):
        os.environ['LITELLM_DEFAULT_MODEL'] = 'claude-2'
        client = GenesisClient()
        self.assertEqual(client.default_model, 'claude-2')
        del os.environ['LITELLM_DEFAULT_MODEL']  # Clean up the environment variable

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