test: update hf test to check if client closed

This commit is contained in:
Krrish Dholakia 2024-12-12 11:34:50 -08:00
parent 90f9aded9f
commit aa7f416b7f

View file

@ -3969,14 +3969,15 @@ def test_completion_hf_prompt_array():
def test_text_completion_stream():
try:
response = text_completion(
model="huggingface/mistralai/Mistral-7B-v0.1",
prompt="good morning",
stream=True,
max_tokens=10,
)
for chunk in response:
print(f"chunk: {chunk}")
for _ in range(2): # check if closed client used
response = text_completion(
model="huggingface/mistralai/Mistral-7B-v0.1",
prompt="good morning",
stream=True,
max_tokens=10,
)
for chunk in response:
print(f"chunk: {chunk}")
except Exception as e:
pytest.fail(f"GOT exception for HF In streaming{e}")