From 5c45e69a5ed811107440bd6990aa7838ce7d40e0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 6 Jan 2024 14:09:10 +0530 Subject: [PATCH] test(test_proxy_server_keys.py): add logic for connecting/disconnecting from http server --- litellm/tests/test_proxy_server_keys.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/litellm/tests/test_proxy_server_keys.py b/litellm/tests/test_proxy_server_keys.py index 73dbf24a31b..4aa2c2e26e5 100644 --- a/litellm/tests/test_proxy_server_keys.py +++ b/litellm/tests/test_proxy_server_keys.py @@ -29,6 +29,7 @@ from litellm.proxy.proxy_server import ( router, save_worker_config, startup_event, + shutdown_event, ) # Replace with the actual module where your FastAPI router is defined filepath = os.path.dirname(os.path.abspath(__file__)) @@ -60,6 +61,11 @@ async def wrapper_startup_event(): await startup_event() +@app.on_event("shutdown") +async def wrapper_shutdown_event(): + await shutdown_event() + + # Here you create a fixture that will be used by your tests # Make sure the fixture returns TestClient(app) @pytest.fixture(autouse=True)