diff --git a/litellm/tests/test_exceptions.py b/litellm/tests/test_exceptions.py index 1cb599206b4..82f52fa015d 100644 --- a/litellm/tests/test_exceptions.py +++ b/litellm/tests/test_exceptions.py @@ -352,6 +352,24 @@ def test_completion_mistral_exception(): # test_completion_mistral_exception() +def test_content_policy_exceptionimage_generation_openai(): + try: + litellm.set_verbose = True + response = litellm.image_generation( + prompt="a very bad prompt", model="dall-e-3" + ) + print(f"response: {response}") + assert len(response.data) > 0 + except litellm.ContentPolicyViolationError as e: + print("caught a content policy violation error! Passed") + pass + except Exception as e: + pytest.fail(f"An exception occurred - {str(e)}") + + +# test_image_generation_openai() + + # # test_invalid_request_error(model="command-nightly") # # Test 3: Rate Limit Errors # def test_model_call(model):