diff options
Diffstat (limited to 'etc/unittest/client.py')
| -rw-r--r-- | etc/unittest/client.py | 10 |
1 files changed, 6 insertions, 4 deletions
diff --git a/etc/unittest/client.py b/etc/unittest/client.py index c63edbd2..ec8aa4b7 100644 --- a/etc/unittest/client.py +++ b/etc/unittest/client.py @@ -35,15 +35,17 @@ class TestPassModel(unittest.TestCase): response = client.chat.completions.create(messages, "Hello", stream=True) for chunk in response: self.assertIsInstance(chunk, ChatCompletionChunk) - self.assertIsInstance(chunk.choices[0].delta.content, str) + 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), 2) + self.assertEqual(len(response), 3) for chunk in response: - self.assertEqual(chunk.choices[0].delta.content, "You ") + if chunk.choices[0].delta.content is not None: + self.assertEqual(chunk.choices[0].delta.content, "You ") - def no_test_stop(self): + 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"]) |