diff options
Diffstat (limited to '')
| -rw-r--r-- | etc/unittest/client.py | 38 |
1 files changed, 20 insertions, 18 deletions
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() |