From cca71b41b35df116c15e57f042a8017db2d36991 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 6 Mar 2025 08:42:17 -0800 Subject: [PATCH] test_e2e_batches_files --- .../test_openai_batches_endpoint.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/openai_misc_endpoints_tests/test_openai_batches_endpoint.py b/tests/openai_misc_endpoints_tests/test_openai_batches_endpoint.py index cf575d62bed..3b6527a11b6 100644 --- a/tests/openai_misc_endpoints_tests/test_openai_batches_endpoint.py +++ b/tests/openai_misc_endpoints_tests/test_openai_batches_endpoint.py @@ -101,7 +101,7 @@ def await_batch_completion(batch_id: str, custom_llm_provider: str): ) if batch.status == "completed": print(f"Batch {batch_id} completed.") - return + return batch.id tries += 1 print(f"waiting for batch to complete... (attempt {tries}/{max_tries})") @@ -173,9 +173,11 @@ def test_e2e_batches_files(custom_llm_provider): # azure takes very long to complete a batch return else: - await_batch_completion( + response_batch_id = await_batch_completion( batch_id=batch_id, custom_llm_provider=custom_llm_provider ) + if response_batch_id is None: + return write_content_to_file( batch_id=batch_id,