fix more tests

upstash · Jul 4, 2024 · 982561a · 982561a
1 parent 2007399
commit 982561a
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/tests/asyncio/test_chat.py b/tests/asyncio/test_chat.py
@@ -178,7 +178,10 @@ async def test_chat_custom_provider_streaming_async(async_qstash: AsyncQStash) -
         if i == 0:
             assert r.choices[0].delta.role is not None
         else:
-            assert r.choices[0].delta.content is not None
+            assert (
+                r.choices[0].delta.content is not None
+                or r.choices[0].finish_reason is not None
+            )
 
         i += 1
 
@@ -215,6 +218,9 @@ async def test_prompt_custom_provider_streaming_async(
         if i == 0:
             assert r.choices[0].delta.role is not None
         else:
-            assert r.choices[0].delta.content is not None
+            assert (
+                r.choices[0].delta.content is not None
+                or r.choices[0].finish_reason is not None
+            )
 
         i += 1
diff --git a/tests/test_chat.py b/tests/test_chat.py
@@ -151,7 +151,10 @@ def test_chat_custom_provider_streaming(qstash: QStash) -> None:
         if i == 0:
             assert r.choices[0].delta.role is not None
         else:
-            assert r.choices[0].delta.content is not None
+            assert (
+                r.choices[0].delta.content is not None
+                or r.choices[0].finish_reason is not None
+            )
 
 
 def test_prompt_custom_provider(qstash: QStash) -> None:
@@ -181,4 +184,7 @@ def test_prompt_custom_provider_streaming(qstash: QStash) -> None:
         if i == 0:
             assert r.choices[0].delta.role is not None
         else:
-            assert r.choices[0].delta.content is not None
+            assert (
+                r.choices[0].delta.content is not None
+                or r.choices[0].finish_reason is not None
+            )