From 18c93d99cd5bf4425948bdaeb4907b15fd2c3af8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 21 Jun 2025 14:17:57 -0700 Subject: [PATCH] test ui sso --- .../proxy/management_endpoints/test_ui_sso.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 8430fa6ee3..866a63e1ce 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -444,6 +444,7 @@ async def test_default_team_params(team_params): mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object) @@ -492,6 +493,7 @@ async def test_create_team_without_default_params(): mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object) @@ -521,7 +523,7 @@ async def test_create_team_without_default_params(): def test_apply_user_info_values_to_sso_user_defined_values(): - from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy._types import LiteLLM_UserTable, SSOUserDefinedValues from litellm.proxy.management_endpoints.ui_sso import ( apply_user_info_values_to_sso_user_defined_values, ) @@ -532,10 +534,13 @@ def test_apply_user_info_values_to_sso_user_defined_values(): user_role="admin", ) - user_defined_values = { + user_defined_values: SSOUserDefinedValues = { + "models": [], "user_id": "456", "user_email": "test@example.com", "user_role": "admin", + "max_budget": None, + "budget_duration": None, } sso_user_defined_values = apply_user_info_values_to_sso_user_defined_values( @@ -543,6 +548,7 @@ def test_apply_user_info_values_to_sso_user_defined_values(): user_defined_values=user_defined_values, ) + assert sso_user_defined_values is not None assert sso_user_defined_values["user_id"] == "123" @@ -582,8 +588,8 @@ async def test_get_user_info_from_db(): "user_email": user_email, "user_defined_values": user_defined_values, } - with patch.object( - litellm.proxy.management_endpoints.ui_sso, "get_user_object" + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_object" ) as mock_get_user_object: user_info = await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() @@ -623,8 +629,8 @@ async def test_get_user_info_from_db_alternate_user_id(): "user_defined_values": user_defined_values, "alternate_user_id": "krrishd-email1234", } - with patch.object( - litellm.proxy.management_endpoints.ui_sso, "get_user_object" + with patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_object" ) as mock_get_user_object: user_info = await get_user_info_from_db(**args) mock_get_user_object.assert_called_once()