Add event expectations (#507)

* Add event expectations * Fix
cohere-ai · Jun 10, 2024 · e874421 · e874421
1 parent 2de2ee2
commit e874421
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/tests/test_async_client.py b/tests/test_async_client.py
@@ -49,10 +49,17 @@ async def test_chat_stream(self) -> None:
             connectors=[ChatConnector(id="web-search")]
         )
 
+        events = set()
+
         async for chat_event in stream:
+            events.add(chat_event.event_type)
             if chat_event.event_type == "text-generation":
                 print(chat_event.text)
 
+        self.assertTrue("text-generation" in events)
+        self.assertTrue("stream-start" in events)
+        self.assertTrue("stream-end" in events)
+
     async def test_stream_equals_true(self) -> None:
         with self.assertRaises(ValueError):
             await self.co.chat(

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -47,10 +47,17 @@ def test_chat_stream(self) -> None:
             connectors=[ChatConnector(id="web-search")]
         )
 
+        events = set()
+
         for chat_event in stream:
+            events.add(chat_event.event_type)
             if chat_event.event_type == "text-generation":
                 print(chat_event.text)
 
+        self.assertTrue("text-generation" in events)
+        self.assertTrue("stream-start" in events)
+        self.assertTrue("stream-end" in events)
+
     def test_stream_equals_true(self) -> None:
         with self.assertRaises(ValueError):
             co.chat(