From da4bd47e3ed6d3b9687c9b4fbba1f98f3aba91da Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 13 Jul 2024 15:04:13 -0700
Subject: [PATCH] test: test fixes

---
 litellm/__init__.py         |  1 +
 litellm/exceptions.py       | 25 +++++++++++++++++++++++++
 litellm/main.py             |  2 +-
 litellm/tests/test_rules.py | 15 ++++++++-------
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index 898a6fb9b6..5c6b2400a0 100644
--- a/litellm/__init__.py
+++ b/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
diff --git a/litellm/exceptions.py b/litellm/exceptions.py
index d85510b1d8..414b3e002a 100644
--- a/litellm/exceptions.py
+++ b/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
diff --git a/litellm/main.py b/litellm/main.py
index 0aeff31880..e01603b7e7 100644
--- a/litellm/main.py
+++ b/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
diff --git a/litellm/tests/test_rules.py b/litellm/tests/test_rules.py
index 0bafbf48f7..20ab48723c 100644
--- a/litellm/tests/test_rules.py
+++ b/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