phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test msft sso handler openid

Browse source
This commit is contained in:
Ishaan Jaff 2025-04-09 15:20:30 -07:00
parent 0601ae55c4
commit a1e88a6c49
2 changed files with 3 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -363,7 +363,6 @@ async def auth_callback(request: Request): # noqa: PLR0915
request=request, request=request,
microsoft_client_id=microsoft_client_id, microsoft_client_id=microsoft_client_id,
redirect_url=redirect_url, redirect_url=redirect_url,
jwt_handler=jwt_handler,
) )
elif generic_client_id is not None: elif generic_client_id is not None:
result = await get_generic_sso_response( result = await get_generic_sso_response(
@ -856,7 +855,6 @@ class MicrosoftSSOHandler:
request: Request, request: Request,
microsoft_client_id: str, microsoft_client_id: str,
redirect_url: str, redirect_url: str,
jwt_handler: JWTHandler,
return_raw_sso_response: bool = False, return_raw_sso_response: bool = False,
) -> Union[CustomOpenID, OpenID, dict]: ) -> Union[CustomOpenID, OpenID, dict]:
""" """
@ -905,14 +903,13 @@ class MicrosoftSSOHandler:
result = MicrosoftSSOHandler.openid_from_response( result = MicrosoftSSOHandler.openid_from_response(
response=original_msft_result, response=original_msft_result,
jwt_handler=jwt_handler,
team_ids=user_team_ids, team_ids=user_team_ids,
) )
return result return result
@staticmethod @staticmethod
def openid_from_response( def openid_from_response(
response: Optional[dict], jwt_handler: JWTHandler, team_ids: List[str] response: Optional[dict], team_ids: List[str]
) -> CustomOpenID: ) -> CustomOpenID:
response = response or {} response = response or {}
verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}")
@ -1151,7 +1148,6 @@ async def debug_sso_callback(request: Request):
request=request, request=request,
microsoft_client_id=microsoft_client_id, microsoft_client_id=microsoft_client_id,
redirect_url=redirect_url, redirect_url=redirect_url,
jwt_handler=jwt_handler,
return_raw_sso_response=True, return_raw_sso_response=True,
) )
elif generic_client_id is not None: elif generic_client_id is not None:

24
tests/litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -32,23 +32,14 @@ def test_microsoft_sso_handler_openid_from_response():
"surname": "User", "surname": "User",
"some_other_field": "value", "some_other_field": "value",
} }
# Create a mock JWTHandler that returns predetermined team IDs
mock_jwt_handler = MagicMock(spec=JWTHandler)
expected_team_ids = ["team1", "team2"] expected_team_ids = ["team1", "team2"]
mock_jwt_handler.get_team_ids_from_jwt.return_value = expected_team_ids
# Act # Act
# Call the method being tested # Call the method being tested
result = MicrosoftSSOHandler.openid_from_response( result = MicrosoftSSOHandler.openid_from_response(
response=mock_response, jwt_handler=mock_jwt_handler response=mock_response, team_ids=expected_team_ids
) )
# Assert # Assert
# Verify the JWT handler was called with the correct parameters
mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with(
cast(dict, mock_response)
)
# Check that the result is a CustomOpenID object with the expected values # Check that the result is a CustomOpenID object with the expected values
assert isinstance(result, CustomOpenID) assert isinstance(result, CustomOpenID)
@ -64,13 +55,9 @@ def test_microsoft_sso_handler_openid_from_response():
def test_microsoft_sso_handler_with_empty_response(): def test_microsoft_sso_handler_with_empty_response():
# Arrange # Arrange
# Test with None response # Test with None response
mock_jwt_handler = MagicMock(spec=JWTHandler)
mock_jwt_handler.get_team_ids_from_jwt.return_value = []
# Act # Act
result = MicrosoftSSOHandler.openid_from_response( result = MicrosoftSSOHandler.openid_from_response(response=None, team_ids=[])
response=None, jwt_handler=mock_jwt_handler
)
# Assert # Assert
assert isinstance(result, CustomOpenID) assert isinstance(result, CustomOpenID)
@ -82,14 +69,10 @@ def test_microsoft_sso_handler_with_empty_response():
assert result.last_name is None assert result.last_name is None
assert result.team_ids == [] assert result.team_ids == []
# Make sure the JWT handler was called with an empty dict
mock_jwt_handler.get_team_ids_from_jwt.assert_called_once_with({})
def test_get_microsoft_callback_response(): def test_get_microsoft_callback_response():
# Arrange # Arrange
mock_request = MagicMock(spec=Request) mock_request = MagicMock(spec=Request)
mock_jwt_handler = MagicMock(spec=JWTHandler)
mock_response = { mock_response = {
"mail": "microsoft_user@example.com", "mail": "microsoft_user@example.com",
"displayName": "Microsoft User", "displayName": "Microsoft User",
@ -115,7 +98,6 @@ def test_get_microsoft_callback_response():
request=mock_request, request=mock_request,
microsoft_client_id="mock_client_id", microsoft_client_id="mock_client_id",
redirect_url="http://mock_redirect_url", redirect_url="http://mock_redirect_url",
jwt_handler=mock_jwt_handler,
) )
) )
@ -132,7 +114,6 @@ def test_get_microsoft_callback_response():
def test_get_microsoft_callback_response_raw_sso_response(): def test_get_microsoft_callback_response_raw_sso_response():
# Arrange # Arrange
mock_request = MagicMock(spec=Request) mock_request = MagicMock(spec=Request)
mock_jwt_handler = MagicMock(spec=JWTHandler)
mock_response = { mock_response = {
"mail": "microsoft_user@example.com", "mail": "microsoft_user@example.com",
"displayName": "Microsoft User", "displayName": "Microsoft User",
@ -157,7 +138,6 @@ def test_get_microsoft_callback_response_raw_sso_response():
request=mock_request, request=mock_request,
microsoft_client_id="mock_client_id", microsoft_client_id="mock_client_id",
redirect_url="http://mock_redirect_url", redirect_url="http://mock_redirect_url",
jwt_handler=mock_jwt_handler,
return_raw_sso_response=True, return_raw_sso_response=True,
) )
) )