test: test fixes

litellm/__init__.py

@@ -876,6 +876,7 @@ from .exceptions import (
     InternalServerError,
     JSONSchemaValidationError,
     LITELLM_EXCEPTION_TYPES,
+    MockException,
 )
 from .budget_manager import BudgetManager
 from .proxy.proxy_cli import run_server

litellm/exceptions.py

@@ -723,3 +723,28 @@ class InvalidRequestError(openai.BadRequestError):  # type: ignore
         super().__init__(
             self.message, f"{self.model}"
         )  # Call the base class constructor with the parameters it needs
+
+
+class MockException(openai.APIError):
+    # used for testing
+    def __init__(
+        self,
+        status_code,
+        message,
+        llm_provider,
+        model,
+        request: Optional[httpx.Request] = None,
+        litellm_debug_info: Optional[str] = None,
+        max_retries: Optional[int] = None,
+        num_retries: Optional[int] = None,
+    ):
+        self.status_code = status_code
+        self.message = "litellm.MockException: {}".format(message)
+        self.llm_provider = llm_provider
+        self.model = model
+        self.litellm_debug_info = litellm_debug_info
+        self.max_retries = max_retries
+        self.num_retries = num_retries
+        if request is None:
+            request = httpx.Request(method="POST", url="https://api.openai.com/v1")
+        super().__init__(self.message, request=request, body=None)  # type: ignore

litellm/main.py

@@ -479,7 +479,7 @@ def mock_completion(
         if isinstance(mock_response, Exception):
             if isinstance(mock_response, openai.APIError):
                 raise mock_response
-            raise litellm.APIError(
+            raise litellm.MockException(
                 status_code=getattr(mock_response, "status_code", 500),  # type: ignore
                 message=getattr(mock_response, "text", str(mock_response)),
                 llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"),  # type: ignore

litellm/tests/test_rules.py

@@ -1,14 +1,18 @@
 #### What this tests ####
 #    This tests setting rules before / after making llm api calls
-import sys, os, time
-import traceback, asyncio
+import asyncio
+import os
+import sys
+import time
+import traceback
+
 import pytest
 
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import litellm
-from litellm import completion, acompletion
+from litellm import acompletion, completion
 
 
 def my_pre_call_rule(input: str):
@@ -126,10 +130,7 @@ def test_post_call_rule_streaming():
         print("Got exception", e)
         print(type(e))
         print(vars(e))
-        assert (
-            "OpenAIException - This violates LiteLLM Proxy Rules. Response too short"
-            in e.message
-        )
+        assert "This violates LiteLLM Proxy Rules. Response too short" in e.message
 
 
 @pytest.mark.asyncio