From 6d8c17967f1121994329c100742a6648891d1eb3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 26 Nov 2024 12:21:55 +0530 Subject: [PATCH] test(test_proxy_server.py): fix team member add tests --- tests/proxy_unit_tests/test_proxy_server.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index d70962858..64bb67b58 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1014,7 +1014,11 @@ async def test_create_team_member_add(prisma_client, new_member_method): with patch( "litellm.proxy.proxy_server.prisma_client.db.litellm_usertable", new_callable=AsyncMock, - ) as mock_litellm_usertable: + ) as mock_litellm_usertable, patch( + "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache", + new=AsyncMock(return_value=team_obj), + ) as mock_team_obj: + mock_client = AsyncMock( return_value=LiteLLM_UserTable( user_id="1234", max_budget=100, user_email="1234" @@ -1193,7 +1197,10 @@ async def test_create_team_member_add_team_admin( with patch( "litellm.proxy.proxy_server.prisma_client.db.litellm_usertable", new_callable=AsyncMock, - ) as mock_litellm_usertable: + ) as mock_litellm_usertable, patch( + "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache", + new=AsyncMock(return_value=team_obj), + ) as mock_team_obj: mock_client = AsyncMock( return_value=LiteLLM_UserTable( user_id="1234", max_budget=100, user_email="1234"