undo merge

Author: alexander-alderman-webb

tests/integrations/litellm/test_litellm.py

@@ -465,9 +465,7 @@ def test_embeddings_no_pii(
         assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"]
 
 
-def test_exception_handling(
-    reset_litellm_executor, sentry_init, capture_events, get_rate_limit_model_response
-):
+def test_exception_handling(sentry_init, capture_events):
     sentry_init(
         integrations=[LiteLLMIntegration()],
         traces_sample_rate=1.0,
@@ -476,24 +474,19 @@ def test_exception_handling(
 
     messages = [{"role": "user", "content": "Hello!"}]
 
-    client = OpenAI(api_key="z")
-
-    model_response = get_rate_limit_model_response()
-
-    with mock.patch.object(
-        client.embeddings._client._client,
-        "send",
-        return_value=model_response,
-    ):
-        with start_transaction(name="litellm test"):
-            with pytest.raises(litellm.RateLimitError):
-                litellm.completion(
-                    model="gpt-3.5-turbo",
-                    messages=messages,
-                    client=client,
-                )
+    with start_transaction(name="litellm test"):
+        kwargs = {
+            "model": "gpt-3.5-turbo",
+            "messages": messages,
+        }
 
-            litellm_utils.executor.shutdown(wait=True)
+        _input_callback(kwargs)
+        _failure_callback(
+            kwargs,
+            Exception("API rate limit reached"),
+            datetime.now(),
+            datetime.now(),
+        )
 
     # Should have error event and transaction
     assert len(events) >= 1