From 6ec058711ab8e77edb674cf1fd7414939c5014e5 Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Thu, 2 May 2024 12:41:30 -0700 Subject: [PATCH] Make unnecessary to pass extra arg for mock object Modify `mock_patch_acompletion` to be a context manager instead of a function that returns a mock object. This way, the mock object is created and yielded by the context manager, and the test function doesn't need to pass the mock object as an argument. --- litellm/tests/test_proxy_server.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 3c192ca22..bdf1d893c 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -2,6 +2,7 @@ import sys, os import traceback from unittest import mock from dotenv import load_dotenv +import contextlib load_dotenv() import os, io @@ -37,6 +38,7 @@ token = "sk-1234" headers = {"Authorization": f"Bearer {token}"} +@contextlib.contextmanager def mock_patch_acompletion(): async def side_effect(*args, **kwargs): return { @@ -50,10 +52,11 @@ def mock_patch_acompletion(): ], } - return mock.patch( + with mock.patch( "litellm.proxy.proxy_server.llm_router.acompletion", side_effect=side_effect, - ) + ): + yield @pytest.fixture(scope="function") @@ -125,7 +128,7 @@ def test_chat_completion_azure(client_no_auth): @mock_patch_acompletion() -def test_openai_deployments_model_chat_completions_azure(_mock_acompletion, client_no_auth): +def test_openai_deployments_model_chat_completions_azure(client_no_auth): global headers try: # Your test data