[fix] change stream tests

edenai · Nov 20, 2023 · bd8f68d · bd8f68d
1 parent 807b571
commit bd8f68d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/edenai_apis/features/text/chat/__init__.py b/edenai_apis/features/text/chat/__init__.py
@@ -1,2 +1,2 @@
 from .chat_args import chat_arguments
-from .chat_dataclass import ChatDataClass, ChatMessageDataClass, StreamChat
+from .chat_dataclass import ChatDataClass, ChatMessageDataClass, StreamChat, ChatStreamResponse
diff --git a/edenai_apis/tests/features/test_text_stream.py b/edenai_apis/tests/features/test_text_stream.py
@@ -4,7 +4,7 @@
 import pytest
 
 from edenai_apis import Text
-from edenai_apis.features.text.chat import StreamChat
+from edenai_apis.features.text.chat import StreamChat, ChatStreamResponse
 from edenai_apis.interface import list_providers
 from edenai_apis.loaders.data_loader import FeatureDataEnum
 from edenai_apis.loaders.loaders import load_feature
@@ -38,4 +38,4 @@ def test_stream(self, provider):
         assert isinstance(chat_output.standardized_response.stream, Iterator)
 
         for chunk in chat_output.standardized_response.stream:
-            assert isinstance(chunk, str)
+            assert isinstance(chunk, ChatStreamResponse)