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