From 982561acfabb1b4521b3c10ffdc88b92e18bd905 Mon Sep 17 00:00:00 2001
From: Metin Dumandag <29387993+mdumandag@users.noreply.github.com>
Date: Thu, 4 Jul 2024 13:34:48 +0300
Subject: [PATCH] fix more tests

---
 tests/asyncio/test_chat.py | 10 ++++++++--
 tests/test_chat.py         | 10 ++++++++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/tests/asyncio/test_chat.py b/tests/asyncio/test_chat.py
index 9ab1e6a..56e51c1 100644
--- 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
index 6b9b7fb..da0ca50 100644
--- 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
+            )