Test_cases.py

import unittest
from unittest.mock import MagicMock

class TestModel(unittest.TestCase):
    def test_init(self):
        model = MockModel("api_key", "mock_model")
        self.assertEqual(model.api_key, "api_key")
        self.assertEqual(model.model, "mock_model")

    def test_supported_models(self):
        self.assertEqual(MockModel.supported_models(), ["mock_model"])

    def test_retry_decorator(self):
        model = MockModel("api_key", "mock_model")
        model.api_wait = 2
        model.api_retry = 3
        decorator = model._retry_decorator()
        self.assertIsNotNone(decorator)

    def test_execute_with_retry(self):
        model = MockModel("api_key", "mock_model")
        model.api_wait = 2
        model.api_retry = 3
        responses = model.execute_with_retry(["prompt"])
        self.assertEqual(responses, ["response"])


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