From 7934610f632ca2e59a50eec2ddb958e82b57256f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 10 May 2025 16:15:48 -0700
Subject: [PATCH] test: update test to handle rate limit error

---
 tests/image_gen_tests/test_image_generation.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py
index 21f96095a7..885df94e19 100644
--- a/tests/image_gen_tests/test_image_generation.py
+++ b/tests/image_gen_tests/test_image_generation.py
@@ -207,6 +207,8 @@ def test_image_generation_azure_dall_e_3():
         pass  # OpenAI randomly raises these errors - skip when they occur
     except litellm.InternalServerError:
         pass
+    except litellm.RateLimitError as e:
+        pass
     except Exception as e:
         if "Your task failed as a result of our safety system." in str(e):
             pass