From d3bbf5bdff39e3a17eec5d477cfff4cfd71d895f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 11 Dec 2023 21:32:53 -0800
Subject: [PATCH] test(test_proxy_server_keys.py): use global client

---
 litellm/tests/langfuse.log              |  3 +++
 litellm/tests/test_proxy_server_keys.py | 31 ++++++++++++-------------
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/litellm/tests/langfuse.log b/litellm/tests/langfuse.log
index e69de29bb..c78eb0911 100644
--- a/litellm/tests/langfuse.log
+++ b/litellm/tests/langfuse.log
@@ -0,0 +1,3 @@
+Giving up get_data(...) after 3 tries (prisma.errors.ClientNotConnectedError: Client is not connected to the query engine, you must call `connect()` before attempting to query data.)
+Giving up get_data(...) after 3 tries (prisma.errors.ClientNotConnectedError: Client is not connected to the query engine, you must call `connect()` before attempting to query data.)
+Giving up get_data(...) after 3 tries (prisma.errors.ClientNotConnectedError: Client is not connected to the query engine, you must call `connect()` before attempting to query data.)
diff --git a/litellm/tests/test_proxy_server_keys.py b/litellm/tests/test_proxy_server_keys.py
index f6a0599e6..db083c30c 100644
--- a/litellm/tests/test_proxy_server_keys.py
+++ b/litellm/tests/test_proxy_server_keys.py
@@ -1,4 +1,4 @@
-import sys, os, time, asyncio
+import sys, os, time
 import traceback
 from dotenv import load_dotenv
 
@@ -23,22 +23,22 @@ from concurrent.futures import ThreadPoolExecutor
 # test /chat/completion request to the proxy
 from fastapi.testclient import TestClient
 from fastapi import FastAPI
-from litellm.proxy.proxy_server import router, save_worker_config, initialize, startup_event  # Replace with the actual module where your FastAPI router is defined
+from litellm.proxy.proxy_server import router, save_worker_config, startup_event  # Replace with the actual module where your FastAPI router is defined
+filepath = os.path.dirname(os.path.abspath(__file__))
+config_fp = f"{filepath}/test_configs/test_config.yaml"
+save_worker_config(config=config_fp, model=None, alias=None, api_base=None, api_version=None, debug=False, temperature=None, max_tokens=None, request_timeout=600, max_budget=None, telemetry=False, drop_params=True, add_function_to_prompt=False, headers=None, save=False, use_queue=False)
+app = FastAPI()
+app.include_router(router)  # Include your router in the test app
+@app.on_event("startup")
+async def wrapper_startup_event():
+    await startup_event()
 
 # Here you create a fixture that will be used by your tests
 # Make sure the fixture returns TestClient(app)
-
-@pytest.fixture(scope="module")
+@pytest.fixture(autouse=True)
 def client():
-    filepath = os.path.dirname(os.path.abspath(__file__))
-    config_fp = f"{filepath}/test_configs/test_config.yaml"
-    save_worker_config(config=config_fp, model=None, alias=None, api_base=None, api_version=None, debug=False, temperature=None, max_tokens=None, request_timeout=600, max_budget=None, telemetry=False, drop_params=True, add_function_to_prompt=False, headers=None, save=False, use_queue=False)
-    initialize(config=config_fp)
-    app = FastAPI()
-    app.include_router(router)  # Include your router in the test app
-    import asyncio
-    asyncio.run(startup_event())
-    return TestClient(app)
+    with TestClient(app) as client:
+        yield client
 
 def test_add_new_key(client):
     try:
@@ -71,8 +71,8 @@ def test_add_new_key(client):
 
 # # Run the test - only runs via pytest
 
-@pytest.mark.asyncio
-async def test_add_new_key_max_parallel_limit(client):
+
+def test_add_new_key_max_parallel_limit(client):
     try:
         # Your test data
         test_data = {"duration": "20m", "max_parallel_requests": 1}
@@ -88,7 +88,6 @@ async def test_add_new_key_max_parallel_limit(client):
         result = response.json()
         def _post_data():
             json_data = {'model': 'azure-model', "messages": [{"role": "user", "content": f"this is a test request, write a short poem {time.time()}"}]}
-            print(f"bearer token key: {result['key']}")
             response = client.post("/chat/completions", json=json_data, headers={"Authorization": f"Bearer {result['key']}"})
             return response
         def _run_in_parallel():