From b4f7771ed132b312f1bff223e80a92116c28c1e5 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 14 Nov 2024 19:02:57 +0530
Subject: [PATCH] test: update test to handle overloaded error

---
 tests/local_testing/test_streaming.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py
index db84234c7..930ef82bd 100644
--- a/tests/local_testing/test_streaming.py
+++ b/tests/local_testing/test_streaming.py
@@ -3333,7 +3333,9 @@ async def test_acompletion_function_call_with_streaming(model):
                 validate_final_streaming_function_calling_chunk(chunk=chunk)
             idx += 1
         # raise Exception("it worked! ")
-    except litellm.InternalServerError or litellm.ServiceUnavailableError:
+    except litellm.InternalServerError:
+        pass
+    except litellm.ServiceUnavailableError:
         pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")