From aa7f416b7f8ac2c0303dfb5963c04e9811bf003b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 12 Dec 2024 11:34:50 -0800 Subject: [PATCH] test: update hf test to check if client closed --- tests/local_testing/test_text_completion.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/local_testing/test_text_completion.py b/tests/local_testing/test_text_completion.py index 19588a9720..cf78b4b3d6 100644 --- a/tests/local_testing/test_text_completion.py +++ b/tests/local_testing/test_text_completion.py @@ -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}")