From 2dc411fdb3aef662d2b86646baac7088c34e4a7b Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 16 Nov 2023 12:24:29 -0800
Subject: [PATCH] (test) hf streaming

---
 litellm/tests/test_text_completion.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/litellm/tests/test_text_completion.py b/litellm/tests/test_text_completion.py
index c16f58d19c7..90505d911a3 100644
--- a/litellm/tests/test_text_completion.py
+++ b/litellm/tests/test_text_completion.py
@@ -155,11 +155,16 @@ def test_completion_hf_prompt_array():
 # test_completion_hf_prompt_array()
 
 def test_text_completion_stream():
-    response = text_completion(
-            model="huggingface/mistralai/Mistral-7B-v0.1", 
-            prompt="good morning",
-            stream=True
-        )
-    for chunk in response:
-        print(chunk)
+    try:
+        response = text_completion(
+                model="huggingface/mistralai/Mistral-7B-v0.1", 
+                prompt="good morning",
+                stream=True,
+                max_tokens=10,
+            )
+        for chunk in response:
+            print(chunk)
+    except Exception as e:
+        pytest.fail(f"GOT exception for HF In streaming{e}")
+
 test_text_completion_stream()