diff --git a/litellm/tests/test_proxy_custom_logger.py b/litellm/tests/test_proxy_custom_logger.py index c8064502e3..6ddc9caacb 100644 --- a/litellm/tests/test_proxy_custom_logger.py +++ b/litellm/tests/test_proxy_custom_logger.py @@ -29,9 +29,10 @@ python_file_path = f"{filepath}/test_configs/custom_callbacks.py" # Use the app fixture in your client fixture -def get_client(config_fp): +@pytest.fixture +def client(): filepath = os.path.dirname(os.path.abspath(__file__)) - config_fp = f"{filepath}/test_configs/{config_fp}" + config_fp = f"{filepath}/test_configs/test_custom_logger.yaml" initialize(config=config_fp) app = FastAPI() app.include_router(router) # Include your router in the test app @@ -49,9 +50,8 @@ headers = { print("Testing proxy custom logger") -def test_embedding(): +def test_embedding(client): try: - client = get_client(config_fp="test_custom_logger.yaml") litellm.set_verbose=False from litellm.proxy.utils import get_instance_fn my_custom_logger = get_instance_fn( @@ -96,10 +96,10 @@ def test_embedding(): pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}") -def test_chat_completion(): +def test_chat_completion(client): try: # Your test data - client = get_client(config_fp="test_custom_logger.yaml") + print("initialized proxy") litellm.set_verbose=False from litellm.proxy.utils import get_instance_fn @@ -160,10 +160,9 @@ def test_chat_completion(): pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}") -def test_chat_completion_stream(): +def test_chat_completion_stream(client): try: # Your test data - client = get_client(config_fp="test_custom_logger.yaml") litellm.set_verbose=False from litellm.proxy.utils import get_instance_fn my_custom_logger = get_instance_fn( diff --git a/litellm/tests/test_proxy_server_keys.py b/litellm/tests/test_proxy_server_keys.py index db083c30c7..eb40c2256e 100644 --- a/litellm/tests/test_proxy_server_keys.py +++ b/litellm/tests/test_proxy_server_keys.py @@ -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, 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() +from litellm.proxy.proxy_server import router, save_worker_config, initialize, startup_event # Replace with the actual module where your FastAPI router is defined # Here you create a fixture that will be used by your tests # Make sure the fixture returns TestClient(app) -@pytest.fixture(autouse=True) + +@pytest.fixture(scope="module") def client(): - with TestClient(app) as client: - yield 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) def test_add_new_key(client): try: