Check if text is not none

cohere-ai · Aug 21, 2023 · 3e35025 · 3e35025
1 parent 0bd0e2c
commit 3e35025
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git 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
@@ -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])