From 6652227c2526faa42d1b5702d7affccb36e93b4d Mon Sep 17 00:00:00 2001
From: yujonglee <yujonglee.dev@gmail.com>
Date: Tue, 4 Jun 2024 13:56:28 +0900
Subject: [PATCH] simple test

---
 litellm/tests/test_opentelemetry.py | 26 ++++++++++++++++++++++++++
 litellm/tests/test_proxy_server.py  | 27 ---------------------------
 2 files changed, 26 insertions(+), 27 deletions(-)
 create mode 100644 litellm/tests/test_opentelemetry.py

diff --git a/litellm/tests/test_opentelemetry.py b/litellm/tests/test_opentelemetry.py
new file mode 100644
index 0000000000..b8d3379bea
--- /dev/null
+++ b/litellm/tests/test_opentelemetry.py
@@ -0,0 +1,26 @@
+import asyncio
+import litellm
+
+from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+
+def test_otel_callback():
+    exporter = InMemorySpanExporter()
+
+    litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter=exporter))]
+
+    litellm.completion(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+    )
+
+    asyncio.run(
+        litellm.acompletion(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "hi"}],
+        )
+    )
+
+    spans = exporter.get_finished_spans()
+    assert len(spans) == 1 + 1
diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py
index 5e8422df3f..114b96872f 100644
--- a/litellm/tests/test_proxy_server.py
+++ b/litellm/tests/test_proxy_server.py
@@ -602,30 +602,3 @@ def test_load_router_config(mock_cache, fake_env_vars):
 
 
 # test_load_router_config()
-
-from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
-
-@mock_patch_acompletion()
-def test_otel_with_proxy_server(mock_acompletion, client_no_auth):
-    exporter = InMemorySpanExporter()
-    litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter=exporter))]
-
-    data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}
-
-    response = client_no_auth.post("/v1/chat/completions", json=data)
-    mock_acompletion.assert_called_once_with(
-        model="gpt-3.5-turbo",
-        messages=[{"role": "user", "content": "hi"}],
-        litellm_call_id=mock.ANY,
-        litellm_logging_obj=mock.ANY,
-        request_timeout=mock.ANY,
-        specific_deployment=True,
-        metadata=mock.ANY,
-        proxy_server_request=mock.ANY,
-    )
-    assert response.status_code == 200
-    assert response.json() == example_completion_result
-
-    spans = exporter.get_finished_spans()
-    assert len(spans) == 0