diff --git a/tests/litellm/proxy/management_endpoints/test_ui_sso.py b/tests/litellm/proxy/management_endpoints/test_ui_sso.py new file mode 100644 index 00000000000..7ad520f7d53 --- /dev/null +++ b/tests/litellm/proxy/management_endpoints/test_ui_sso.py @@ -0,0 +1,81 @@ +import json +import os +import sys +from typing import Optional, cast +from unittest.mock import MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.management_endpoints.types import CustomOpenID +from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler + + +def test_microsoft_sso_handler_openid_from_response(): + # Arrange + # Create a mock response similar to what Microsoft SSO would return + mock_response = { + "mail": "test@example.com", + "displayName": "Test User", + "id": "user123", + "givenName": "Test", + "surname": "User", + "some_other_field": "value", + } + + # Create a mock JWTHandler that returns predetermined team IDs + mock_jwt_handler = MagicMock(spec=JWTHandler) + expected_team_ids = ["team1", "team2"] + mock_jwt_handler.get_team_ids_from_jwt.return_value = expected_team_ids + + # Act + # Call the method being tested + result = MicrosoftSSOHandler.openid_from_response( + response=mock_response, jwt_handler=mock_jwt_handler + ) + + # 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 + assert isinstance(result, CustomOpenID) + assert result.email == "test@example.com" + assert result.display_name == "Test User" + assert result.provider == "microsoft" + assert result.id == "user123" + assert result.first_name == "Test" + assert result.last_name == "User" + assert result.team_ids == expected_team_ids + + +def test_microsoft_sso_handler_with_empty_response(): + # Arrange + # Test with None response + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + # Act + result = MicrosoftSSOHandler.openid_from_response( + response=None, jwt_handler=mock_jwt_handler + ) + + # Assert + assert isinstance(result, CustomOpenID) + assert result.email is None + assert result.display_name is None + assert result.provider == "microsoft" + assert result.id is None + assert result.first_name is None + assert result.last_name is None + 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({})