From c0b56b65758511e7e66c9e12ab3cb019bcdc4ee0 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 9 Jan 2024 17:04:04 +0530
Subject: [PATCH] (test) catch litellm.ContentPolicyViolationError

---
 litellm/tests/test_image_generation.py | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/litellm/tests/test_image_generation.py b/litellm/tests/test_image_generation.py
index 973ec29bb8..3c792f8022 100644
--- a/litellm/tests/test_image_generation.py
+++ b/litellm/tests/test_image_generation.py
@@ -19,7 +19,7 @@ import litellm
 
 
 def test_image_generation_openai():
-    try: 
+    try:
         litellm.set_verbose = True
         response = litellm.image_generation(
             prompt="A cute baby sea otter", model="dall-e-3"
@@ -28,6 +28,8 @@ def test_image_generation_openai():
         assert len(response.data) > 0
     except litellm.RateLimitError as e:
         pass
+    except litellm.ContentPolicyViolationError:
+        pass  # OpenAI randomly raises these errors - skip when they occur
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
 
@@ -36,22 +38,27 @@ def test_image_generation_openai():
 
 
 def test_image_generation_azure():
-    try: 
+    try:
         response = litellm.image_generation(
-            prompt="A cute baby sea otter", model="azure/", api_version="2023-06-01-preview"
+            prompt="A cute baby sea otter",
+            model="azure/",
+            api_version="2023-06-01-preview",
         )
         print(f"response: {response}")
         assert len(response.data) > 0
     except litellm.RateLimitError as e:
         pass
+    except litellm.ContentPolicyViolationError:
+        pass  # Azure randomly raises these errors - skip when they occur
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
 
+
 # test_image_generation_azure()
 
 
 def test_image_generation_azure_dall_e_3():
-    try: 
+    try:
         litellm.set_verbose = True
         response = litellm.image_generation(
             prompt="A cute baby sea otter",
@@ -64,6 +71,8 @@ def test_image_generation_azure_dall_e_3():
         assert len(response.data) > 0
     except litellm.RateLimitError as e:
         pass
+    except litellm.ContentPolicyViolationError:
+        pass  # OpenAI randomly raises these errors - skip when they occur
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
 
@@ -71,7 +80,7 @@ def test_image_generation_azure_dall_e_3():
 # test_image_generation_azure_dall_e_3()
 @pytest.mark.asyncio
 async def test_async_image_generation_openai():
-    try: 
+    try:
         response = litellm.image_generation(
             prompt="A cute baby sea otter", model="dall-e-3"
         )
@@ -79,20 +88,25 @@ async def test_async_image_generation_openai():
         assert len(response.data) > 0
     except litellm.RateLimitError as e:
         pass
+    except litellm.ContentPolicyViolationError:
+        pass  # openai randomly raises these errors - skip when they occur
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
 
+
 # asyncio.run(test_async_image_generation_openai())
 
 
 @pytest.mark.asyncio
 async def test_async_image_generation_azure():
-    try: 
+    try:
         response = await litellm.aimage_generation(
             prompt="A cute baby sea otter", model="azure/dall-e-3-test"
         )
         print(f"response: {response}")
     except litellm.RateLimitError as e:
         pass
+    except litellm.ContentPolicyViolationError:
+        pass  # Azure randomly raises these errors - skip when they occur
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")