diff --git a/tests/async/test_async_chat.py b/tests/async/test_async_chat.py index db3ac538c..b24a27d57 100644 --- a/tests/async/test_async_chat.py +++ b/tests/async/test_async_chat.py @@ -36,12 +36,13 @@ async def test_async_chat_stream(async_client): expected_index = 0 expected_text = "" async for token in res: - assert isinstance(token.text, str) - assert len(token.text) > 0 - assert token.index == expected_index + if token.text: + assert isinstance(token.text, str) + assert len(token.text) > 0 + assert token.index == expected_index + expected_text += token.text expected_index += 1 - expected_text += token.text assert res.texts == [expected_text] assert res.conversation_id is not None diff --git a/tests/sync/test_chat.py b/tests/sync/test_chat.py index 931b82155..30843ba2f 100644 --- a/tests/sync/test_chat.py +++ b/tests/sync/test_chat.py @@ -91,13 +91,14 @@ def test_stream(self): expected_index = 0 expected_text = "" for token in prediction: - self.assertIsInstance(token.text, str) - self.assertGreater(len(token.text), 0) + if token.text: + self.assertIsInstance(token.text, str) + self.assertGreater(len(token.text), 0) - self.assertIsInstance(token.index, int) - self.assertEqual(token.index, expected_index) + self.assertIsInstance(token.index, int) + self.assertEqual(token.index, expected_index) - expected_text += token.text + expected_text += token.text expected_index += 1 self.assertEqual(prediction.texts, [expected_text])