diff options
Diffstat (limited to 'etc/unittest')
| -rw-r--r-- | etc/unittest/client.py | 8 |
1 files changed, 5 insertions, 3 deletions
diff --git a/etc/unittest/client.py b/etc/unittest/client.py index 2bc00c2e..ec8aa4b7 100644 --- a/etc/unittest/client.py +++ b/etc/unittest/client.py @@ -35,13 +35,15 @@ 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 test_stop(self): client = Client(provider=YieldProviderMock) |