6 files changed, 157 insertions, 58 deletions

diff --git a/etc/unittest/__main__.py b/etc/unittest/__main__.py
new file mode 100644
index 00000000..243c56b2
--- /dev/null
+++ b/etc/unittest/__main__.py
@@ -0,0 +1,6 @@
+import unittest
+from .asyncio import *
+from .backend import *
+from .main import *
+
+unittest.main()
\ No newline at end of file
diff --git a/etc/unittest/asyncio.py b/etc/unittest/asyncio.py
new file mode 100644
index 00000000..74e29986
--- /dev/null
+++ b/etc/unittest/asyncio.py
@@ -0,0 +1,57 @@
+from .include import DEFAULT_MESSAGES
+import asyncio
+import nest_asyncio
+import unittest
+import g4f
+from g4f import ChatCompletion
+from .mocks import ProviderMock, AsyncProviderMock, AsyncGeneratorProviderMock
+        
+class TestChatCompletion(unittest.TestCase):
+    
+    async def run_exception(self):
+        return ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
+        
+    def test_exception(self):
+        self.assertRaises(g4f.errors.NestAsyncioError, asyncio.run, self.run_exception())
+
+    def test_create(self):
+        result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
+        self.assertEqual("Mock",result)
+        
+    def test_create_generator(self):
+        result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncGeneratorProviderMock)
+        self.assertEqual("Mock",result)
+
+class TestChatCompletionAsync(unittest.IsolatedAsyncioTestCase):
+    
+    async def test_base(self):
+        result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, ProviderMock)
+        self.assertEqual("Mock",result)
+        
+    async def test_async(self):
+        result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
+        self.assertEqual("Mock",result)
+        
+    async def test_create_generator(self):
+        result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, AsyncGeneratorProviderMock)
+        self.assertEqual("Mock",result)
+        
+class TestChatCompletionNestAsync(unittest.IsolatedAsyncioTestCase):
+        
+    def setUp(self) -> None:
+        nest_asyncio.apply()
+        
+    async def test_create(self):
+        result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, ProviderMock)
+        self.assertEqual("Mock",result)
+        
+    async def test_nested(self):
+        result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
+        self.assertEqual("Mock",result)
+        
+    async def test_nested_generator(self):
+        result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncGeneratorProviderMock)
+        self.assertEqual("Mock",result)
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/etc/unittest/backend.py b/etc/unittest/backend.py
new file mode 100644
index 00000000..f5961e2d
--- /dev/null
+++ b/etc/unittest/backend.py
@@ -0,0 +1,38 @@
+from . import include
+import unittest
+from unittest.mock import MagicMock
+from .mocks import ProviderMock
+import g4f
+from g4f.gui.server.backend import Backend_Api, get_error_message
+        
+class TestBackendApi(unittest.TestCase):
+
+    def setUp(self):
+        self.app = MagicMock()
+        self.api = Backend_Api(self.app)
+
+    def test_version(self):
+        response = self.api.get_version()
+        self.assertIn("version", response)
+        self.assertIn("latest_version", response)
+        
+    def test_get_models(self):
+        response = self.api.get_models()
+        self.assertIsInstance(response, list)
+        self.assertTrue(len(response) > 0)
+        
+    def test_get_providers(self):
+        response = self.api.get_providers()
+        self.assertIsInstance(response, list)
+        self.assertTrue(len(response) > 0)
+        
+class TestUtilityFunctions(unittest.TestCase):
+
+    def test_get_error_message(self):
+        g4f.debug.last_provider = ProviderMock
+        exception = Exception("Message")
+        result = get_error_message(exception)
+        self.assertEqual("ProviderMock: Exception: Message", result)
+        
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/etc/unittest/include.py b/etc/unittest/include.py
new file mode 100644
index 00000000..e67fd5a7
--- /dev/null
+++ b/etc/unittest/include.py
@@ -0,0 +1,11 @@
+import sys
+import pathlib
+
+sys.path.append(str(pathlib.Path(__file__).parent.parent.parent))
+
+import g4f
+
+g4f.debug.logging = False
+g4f.debug.version_check = False
+
+DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}]
\ No newline at end of file
diff --git a/etc/unittest/main.py b/etc/unittest/main.py
index ad1fe02d..5a220323 100644
--- a/etc/unittest/main.py
+++ b/etc/unittest/main.py
@@ -1,75 +1,37 @@
-import sys
-import pathlib
+from .include import DEFAULT_MESSAGES
 import unittest
-from unittest.mock import MagicMock
-
-sys.path.append(str(pathlib.Path(__file__).parent.parent.parent))
-
+import asyncio
 import g4f
 from g4f import ChatCompletion, get_last_provider
-from g4f.gui.server.backend import Backend_Api, get_error_message
-from g4f.base_provider import BaseProvider
-
-g4f.debug.logging = False
-g4f.debug.version_check = False
-
-class MockProvider(BaseProvider):
-    working = True
-
-    def create_completion(
-        model, messages, stream, **kwargs
-    ):
-        yield "Mock"
-
-    async def create_async(
-        model, messages, **kwargs
-    ):
-        return "Mock"
-
-class TestBackendApi(unittest.TestCase):
-
-    def setUp(self):
-        self.app = MagicMock()
-        self.api = Backend_Api(self.app)
-
-    def test_version(self):
-        response = self.api.get_version()
-        self.assertIn("version", response)
-        self.assertIn("latest_version", response)
+from g4f.Provider import RetryProvider
+from .mocks import ProviderMock
 
 class TestChatCompletion(unittest.TestCase):
 
     def test_create_default(self):
-        messages = [{'role': 'user', 'content': 'Hello'}]
-        result = ChatCompletion.create(g4f.models.default, messages)
+        result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES)
         if "Good" not in result and "Hi" not in result:
             self.assertIn("Hello", result)
-        
-    def test_get_last_provider(self):
-        messages = [{'role': 'user', 'content': 'Hello'}]
-        ChatCompletion.create(g4f.models.default, messages, MockProvider)
-        self.assertEqual(get_last_provider(), MockProvider)
-        
+
     def test_bing_provider(self):
-        messages = [{'role': 'user', 'content': 'Hello'}]
         provider = g4f.Provider.Bing
-        result = ChatCompletion.create(g4f.models.default, messages, provider)
+        result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, provider)
         self.assertIn("Bing", result)
 
-class TestChatCompletionAsync(unittest.IsolatedAsyncioTestCase):
-    
-    async def test_async(self):
-        messages = [{'role': 'user', 'content': 'Hello'}]
-        result = await ChatCompletion.create_async(g4f.models.default, messages, MockProvider)
-        self.assertEqual("Mock", result)
+class TestGetLastProvider(unittest.TestCase):
 
-class TestUtilityFunctions(unittest.TestCase):
-
-    def test_get_error_message(self):
-        g4f.debug.last_provider = g4f.Provider.Bing
-        exception = Exception("Message")
-        result = get_error_message(exception)
-        self.assertEqual("Bing: Exception: Message", result)
+    def test_get_last_provider(self):
+        ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, ProviderMock)
+        self.assertEqual(get_last_provider(), ProviderMock)
+        
+    def test_get_last_provider_retry(self):
+        ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, RetryProvider([ProviderMock]))
+        self.assertEqual(get_last_provider(), ProviderMock)
+        
+    def test_get_last_provider_async(self):
+        coroutine = ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, ProviderMock)
+        asyncio.run(coroutine)
+        self.assertEqual(get_last_provider(), ProviderMock)
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/etc/unittest/mocks.py b/etc/unittest/mocks.py
new file mode 100644
index 00000000..a9505997
--- /dev/null
+++ b/etc/unittest/mocks.py
@@ -0,0 +1,25 @@
+from g4f.Provider.base_provider import AbstractProvider, AsyncProvider, AsyncGeneratorProvider
+
+class ProviderMock(AbstractProvider):
+    working = True
+
+    def create_completion(
+        model, messages, stream, **kwargs
+    ):
+        yield "Mock"
+        
+class AsyncProviderMock(AsyncProvider):
+    working = True
+    
+    async def create_async(
+        model, messages, **kwargs
+    ):
+        return "Mock"
+
+class AsyncGeneratorProviderMock(AsyncGeneratorProvider):
+    working = True
+    
+    async def create_async_generator(
+        model, messages, stream, **kwargs
+    ):
+        yield "Mock"
\ No newline at end of file