diff options
Diffstat (limited to '')
-rw-r--r-- | etc/unittest/__main__.py | 4 | ||||
-rw-r--r-- | etc/unittest/async_client.py | 56 | ||||
-rw-r--r-- | etc/unittest/client.py | 38 |
3 files changed, 22 insertions, 76 deletions
diff --git a/etc/unittest/__main__.py b/etc/unittest/__main__.py index 351c2bb3..ee748917 100644 --- a/etc/unittest/__main__.py +++ b/etc/unittest/__main__.py @@ -4,8 +4,8 @@ from .backend import * from .main import * from .model import * from .client import * -from .async_client import * +from .client import * from .include import * from .integration import * -unittest.main()
\ No newline at end of file +unittest.main() diff --git a/etc/unittest/async_client.py b/etc/unittest/async_client.py deleted file mode 100644 index 06ea5cdb..00000000 --- a/etc/unittest/async_client.py +++ /dev/null @@ -1,56 +0,0 @@ -import unittest - -from g4f.client import AsyncClient, ChatCompletion, ChatCompletionChunk -from .mocks import AsyncGeneratorProviderMock, ModelProviderMock, YieldProviderMock - -DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}] - -class AsyncTestPassModel(unittest.IsolatedAsyncioTestCase): - - async def test_response(self): - client = AsyncClient(provider=AsyncGeneratorProviderMock) - response = await client.chat.completions.create(DEFAULT_MESSAGES, "") - self.assertIsInstance(response, ChatCompletion) - self.assertEqual("Mock", response.choices[0].message.content) - - async def test_pass_model(self): - client = AsyncClient(provider=ModelProviderMock) - response = await client.chat.completions.create(DEFAULT_MESSAGES, "Hello") - self.assertIsInstance(response, ChatCompletion) - self.assertEqual("Hello", response.choices[0].message.content) - - async def test_max_tokens(self): - client = AsyncClient(provider=YieldProviderMock) - messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] - response = await client.chat.completions.create(messages, "Hello", max_tokens=1) - self.assertIsInstance(response, ChatCompletion) - self.assertEqual("How ", response.choices[0].message.content) - response = await client.chat.completions.create(messages, "Hello", max_tokens=2) - self.assertIsInstance(response, ChatCompletion) - self.assertEqual("How are ", response.choices[0].message.content) - - async def test_max_stream(self): - client = AsyncClient(provider=YieldProviderMock) - messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] - response = await client.chat.completions.create(messages, "Hello", stream=True) - async for chunk in response: - self.assertIsInstance(chunk, ChatCompletionChunk) - if chunk.choices[0].delta.content is not None: - self.assertIsInstance(chunk.choices[0].delta.content, str) - messages = [{'role': 'user', 'content': chunk} for chunk in ["You ", "You ", "Other", "?"]] - response = await client.chat.completions.create(messages, "Hello", stream=True, max_tokens=2) - response = [chunk async for chunk in response] - self.assertEqual(len(response), 3) - for chunk in response: - if chunk.choices[0].delta.content is not None: - self.assertEqual(chunk.choices[0].delta.content, "You ") - - async def test_stop(self): - client = AsyncClient(provider=YieldProviderMock) - messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] - response = await client.chat.completions.create(messages, "Hello", stop=["and"]) - self.assertIsInstance(response, ChatCompletion) - self.assertEqual("How are you?", response.choices[0].message.content) - -if __name__ == '__main__': - unittest.main() diff --git a/etc/unittest/client.py b/etc/unittest/client.py index ec8aa4b7..54e2091f 100644 --- a/etc/unittest/client.py +++ b/etc/unittest/client.py @@ -5,52 +5,54 @@ from .mocks import AsyncGeneratorProviderMock, ModelProviderMock, YieldProviderM DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}] -class TestPassModel(unittest.TestCase): +class AsyncTestPassModel(unittest.IsolatedAsyncioTestCase): - def test_response(self): + async def test_response(self): client = Client(provider=AsyncGeneratorProviderMock) - response = client.chat.completions.create(DEFAULT_MESSAGES, "") + response = await client.chat.completions.async_create(DEFAULT_MESSAGES, "") self.assertIsInstance(response, ChatCompletion) self.assertEqual("Mock", response.choices[0].message.content) - def test_pass_model(self): + async def test_pass_model(self): client = Client(provider=ModelProviderMock) - response = client.chat.completions.create(DEFAULT_MESSAGES, "Hello") + response = await client.chat.completions.async_create(DEFAULT_MESSAGES, "Hello") self.assertIsInstance(response, ChatCompletion) self.assertEqual("Hello", response.choices[0].message.content) - def test_max_tokens(self): + async def test_max_tokens(self): client = Client(provider=YieldProviderMock) messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] - response = client.chat.completions.create(messages, "Hello", max_tokens=1) + response = await client.chat.completions.async_create(messages, "Hello", max_tokens=1) self.assertIsInstance(response, ChatCompletion) self.assertEqual("How ", response.choices[0].message.content) - response = client.chat.completions.create(messages, "Hello", max_tokens=2) + response = await client.chat.completions.async_create(messages, "Hello", max_tokens=2) self.assertIsInstance(response, ChatCompletion) self.assertEqual("How are ", response.choices[0].message.content) - def test_max_stream(self): + async def test_max_stream(self): client = Client(provider=YieldProviderMock) messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] - response = client.chat.completions.create(messages, "Hello", stream=True) - for chunk in response: + response = await client.chat.completions.async_create(messages, "Hello", stream=True) + async for chunk in response: self.assertIsInstance(chunk, ChatCompletionChunk) if chunk.choices[0].delta.content is not None: self.assertIsInstance(chunk.choices[0].delta.content, str) messages = [{'role': 'user', 'content': chunk} for chunk in ["You ", "You ", "Other", "?"]] - response = client.chat.completions.create(messages, "Hello", stream=True, max_tokens=2) - response = list(response) - self.assertEqual(len(response), 3) - for chunk in response: + response = await client.chat.completions.async_create(messages, "Hello", stream=True, max_tokens=2) + response_list = [] + async for chunk in response: + response_list.append(chunk) + self.assertEqual(len(response_list), 3) + for chunk in response_list: if chunk.choices[0].delta.content is not None: self.assertEqual(chunk.choices[0].delta.content, "You ") - def test_stop(self): + async def test_stop(self): client = Client(provider=YieldProviderMock) messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] - response = client.chat.completions.create(messages, "Hello", stop=["and"]) + response = await client.chat.completions.async_create(messages, "Hello", stop=["and"]) self.assertIsInstance(response, ChatCompletion) self.assertEqual("How are you?", response.choices[0].message.content) if __name__ == '__main__': - unittest.main()
\ No newline at end of file + unittest.main() |