phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Make unnecessary to pass extra arg for mock object

Browse source 
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.
This commit is contained in:
Marc Abramowitz 2024-05-02 12:41:30 -07:00
parent 4dfadb0cf4
commit 6ec058711a
1 changed files with 6 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/tests/test_proxy_server.py
View file

@ -2,6 +2,7 @@ import sys, os
import traceback import traceback
from unittest import mock from unittest import mock
from dotenv import load_dotenv from dotenv import load_dotenv
import contextlib
load_dotenv() load_dotenv()
import os, io import os, io
@ -37,6 +38,7 @@ token = "sk-1234"
headers = {"Authorization": f"Bearer {token}"} headers = {"Authorization": f"Bearer {token}"}
@contextlib.contextmanager
def mock_patch_acompletion(): def mock_patch_acompletion():
async def side_effect(*args, **kwargs): async def side_effect(*args, **kwargs):
return { return {
@ -50,10 +52,11 @@ def mock_patch_acompletion():
], ],
} }
return mock.patch( with mock.patch(
"litellm.proxy.proxy_server.llm_router.acompletion", "litellm.proxy.proxy_server.llm_router.acompletion",
side_effect=side_effect, side_effect=side_effect,
) ):
yield
@pytest.fixture(scope="function") @pytest.fixture(scope="function")
@ -125,7 +128,7 @@ def test_chat_completion_azure(client_no_auth):
@mock_patch_acompletion() @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 global headers
try: try:
# Your test data # Your test data