main.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. import sys
  2. import pathlib
  3. import unittest
  4. from unittest.mock import MagicMock
  5. sys.path.append(str(pathlib.Path(__file__).parent.parent.parent))
  6. import g4f
  7. from g4f import ChatCompletion, get_last_provider
  8. from g4f.gui.server.backend import Backend_Api, get_error_message
  9. from g4f.base_provider import BaseProvider
  10. g4f.debug.logging = False
  11. g4f.debug.version_check = False
  12. class MockProvider(BaseProvider):
  13. working = True
  14. def create_completion(
  15. model, messages, stream, **kwargs
  16. ):
  17. yield "Mock"
  18. async def create_async(
  19. model, messages, **kwargs
  20. ):
  21. return "Mock"
  22. class TestBackendApi(unittest.TestCase):
  23. def setUp(self):
  24. self.app = MagicMock()
  25. self.api = Backend_Api(self.app)
  26. def test_version(self):
  27. response = self.api.get_version()
  28. self.assertIn("version", response)
  29. self.assertIn("latest_version", response)
  30. class TestChatCompletion(unittest.TestCase):
  31. def test_create(self):
  32. messages = [{'role': 'user', 'content': 'Hello'}]
  33. result = ChatCompletion.create(g4f.models.default, messages)
  34. self.assertTrue("Hello" in result or "Good" in result)
  35. def test_get_last_provider(self):
  36. messages = [{'role': 'user', 'content': 'Hello'}]
  37. ChatCompletion.create(g4f.models.default, messages, MockProvider)
  38. self.assertEqual(get_last_provider(), MockProvider)
  39. def test_bing_provider(self):
  40. messages = [{'role': 'user', 'content': 'Hello'}]
  41. provider = g4f.Provider.Bing
  42. result = ChatCompletion.create(g4f.models.default, messages, provider)
  43. self.assertTrue("Bing" in result)
  44. class TestChatCompletionAsync(unittest.IsolatedAsyncioTestCase):
  45. async def test_async(self):
  46. messages = [{'role': 'user', 'content': 'Hello'}]
  47. result = await ChatCompletion.create_async(g4f.models.default, messages, MockProvider)
  48. self.assertTrue("Mock" in result)
  49. class TestUtilityFunctions(unittest.TestCase):
  50. def test_get_error_message(self):
  51. g4f.debug.last_provider = g4f.Provider.Bing
  52. exception = Exception("Message")
  53. result = get_error_message(exception)
  54. self.assertEqual("Bing: Exception: Message", result)
  55. if __name__ == '__main__':
  56. unittest.main()