add /sso/readiness route
This commit is contained in:
yuneng-jiang
parent
dc79da871e
commit
2d8f5111af
@ -1126,6 +1126,91 @@ async def get_ui_settings(request: Request):
|
||||
}
|
||||
|
||||
|
||||
@router.get(
|
||||
"/sso/readiness",
|
||||
tags=["experimental"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def sso_readiness():
|
||||
"""
|
||||
Health endpoint for checking SSO readiness.
|
||||
Checks if the configured SSO provider has all required environment variables set in memory.
|
||||
"""
|
||||
microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
|
||||
google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
|
||||
generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)
|
||||
|
||||
# Determine which SSO provider is configured
|
||||
configured_provider = None
|
||||
if google_client_id is not None:
|
||||
configured_provider = "google"
|
||||
elif microsoft_client_id is not None:
|
||||
configured_provider = "microsoft"
|
||||
elif generic_client_id is not None:
|
||||
configured_provider = "generic"
|
||||
|
||||
# If no SSO is configured, return healthy (SSO is optional)
|
||||
if configured_provider is None:
|
||||
return {
|
||||
"status": "healthy",
|
||||
"sso_configured": False,
|
||||
"message": "No SSO provider configured",
|
||||
}
|
||||
|
||||
# Check required environment variables for the configured provider
|
||||
missing_vars = []
|
||||
|
||||
if configured_provider == "google":
|
||||
google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
|
||||
if google_client_secret is None:
|
||||
missing_vars.append("GOOGLE_CLIENT_SECRET")
|
||||
|
||||
elif configured_provider == "microsoft":
|
||||
microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
|
||||
microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
|
||||
if microsoft_client_secret is None:
|
||||
missing_vars.append("MICROSOFT_CLIENT_SECRET")
|
||||
if microsoft_tenant is None:
|
||||
missing_vars.append("MICROSOFT_TENANT")
|
||||
|
||||
elif configured_provider == "generic":
|
||||
generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
|
||||
generic_authorization_endpoint = os.getenv(
|
||||
"GENERIC_AUTHORIZATION_ENDPOINT", None
|
||||
)
|
||||
generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
|
||||
generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)
|
||||
if generic_client_secret is None:
|
||||
missing_vars.append("GENERIC_CLIENT_SECRET")
|
||||
if generic_authorization_endpoint is None:
|
||||
missing_vars.append("GENERIC_AUTHORIZATION_ENDPOINT")
|
||||
if generic_token_endpoint is None:
|
||||
missing_vars.append("GENERIC_TOKEN_ENDPOINT")
|
||||
if generic_userinfo_endpoint is None:
|
||||
missing_vars.append("GENERIC_USERINFO_ENDPOINT")
|
||||
|
||||
# If all required variables are present, return healthy
|
||||
if len(missing_vars) == 0:
|
||||
return {
|
||||
"status": "healthy",
|
||||
"sso_configured": True,
|
||||
"provider": configured_provider,
|
||||
"message": f"{configured_provider.capitalize()} SSO is properly configured",
|
||||
}
|
||||
|
||||
# If some variables are missing, return unhealthy
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail={
|
||||
"status": "unhealthy",
|
||||
"sso_configured": True,
|
||||
"provider": configured_provider,
|
||||
"missing_environment_variables": missing_vars,
|
||||
"message": f"{configured_provider.capitalize()} SSO is configured but missing required environment variables: {', '.join(missing_vars)}",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class SSOAuthenticationHandler:
|
||||
"""
|
||||
Handler for SSO Authentication across all SSO providers
|
||||
|
||||
@ -3043,3 +3043,236 @@ class TestAddMissingTeamMember:
|
||||
assert set(added_teams) == set(
|
||||
expected_teams_added
|
||||
), f"Expected teams {expected_teams_added}, but got {added_teams}"
|
||||
|
||||
|
||||
class TestSSOReadinessEndpoint:
|
||||
"""Test the /sso/readiness endpoint"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sso_readiness_no_sso_configured(self):
|
||||
"""Test that readiness returns healthy when no SSO is configured"""
|
||||
from fastapi.testclient import TestClient
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
client = TestClient(app)
|
||||
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
response = client.get("/sso/readiness")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "healthy"
|
||||
assert data["sso_configured"] is False
|
||||
assert data["message"] == "No SSO provider configured"
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sso_readiness_google_fully_configured(self):
|
||||
"""Test that readiness returns healthy when Google SSO is fully configured"""
|
||||
from fastapi.testclient import TestClient
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
client = TestClient(app)
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"GOOGLE_CLIENT_ID": "test-google-client-id",
|
||||
"GOOGLE_CLIENT_SECRET": "test-google-secret",
|
||||
},
|
||||
clear=True,
|
||||
):
|
||||
response = client.get("/sso/readiness")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "healthy"
|
||||
assert data["sso_configured"] is True
|
||||
assert data["provider"] == "google"
|
||||
assert "Google SSO is properly configured" in data["message"]
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sso_readiness_google_missing_secret(self):
|
||||
"""Test that readiness returns unhealthy when Google SSO is missing GOOGLE_CLIENT_SECRET"""
|
||||
from fastapi.testclient import TestClient
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
client = TestClient(app)
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"GOOGLE_CLIENT_ID": "test-google-client-id"},
|
||||
clear=True,
|
||||
):
|
||||
response = client.get("/sso/readiness")
|
||||
|
||||
assert response.status_code == 503
|
||||
data = response.json()["detail"]
|
||||
assert data["status"] == "unhealthy"
|
||||
assert data["sso_configured"] is True
|
||||
assert data["provider"] == "google"
|
||||
assert "GOOGLE_CLIENT_SECRET" in data["missing_environment_variables"]
|
||||
assert "Google SSO is configured but missing required environment variables" in data["message"]
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"env_vars,expected_status,expected_provider,expected_missing_vars",
|
||||
[
|
||||
(
|
||||
{
|
||||
"MICROSOFT_CLIENT_ID": "test-microsoft-client-id",
|
||||
"MICROSOFT_CLIENT_SECRET": "test-microsoft-secret",
|
||||
"MICROSOFT_TENANT": "test-tenant",
|
||||
},
|
||||
200,
|
||||
"microsoft",
|
||||
[],
|
||||
),
|
||||
(
|
||||
{"MICROSOFT_CLIENT_ID": "test-microsoft-client-id"},
|
||||
503,
|
||||
"microsoft",
|
||||
["MICROSOFT_CLIENT_SECRET", "MICROSOFT_TENANT"],
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_sso_readiness_microsoft_configurations(
|
||||
self, env_vars, expected_status, expected_provider, expected_missing_vars
|
||||
):
|
||||
"""Test Microsoft SSO readiness with both fully configured and missing variables"""
|
||||
from fastapi.testclient import TestClient
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
client = TestClient(app)
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=True):
|
||||
response = client.get("/sso/readiness")
|
||||
|
||||
assert response.status_code == expected_status
|
||||
|
||||
if expected_status == 200:
|
||||
data = response.json()
|
||||
assert data["sso_configured"] is True
|
||||
assert data["provider"] == expected_provider
|
||||
assert data["status"] == "healthy"
|
||||
assert "Microsoft SSO is properly configured" in data["message"]
|
||||
else:
|
||||
data = response.json()["detail"]
|
||||
assert data["sso_configured"] is True
|
||||
assert data["provider"] == expected_provider
|
||||
assert data["status"] == "unhealthy"
|
||||
assert set(data["missing_environment_variables"]) == set(
|
||||
expected_missing_vars
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"env_vars,expected_status,expected_provider,expected_missing_vars",
|
||||
[
|
||||
(
|
||||
{
|
||||
"GENERIC_CLIENT_ID": "test-generic-client-id",
|
||||
"GENERIC_CLIENT_SECRET": "test-generic-secret",
|
||||
"GENERIC_AUTHORIZATION_ENDPOINT": "https://auth.example.com/authorize",
|
||||
"GENERIC_TOKEN_ENDPOINT": "https://auth.example.com/token",
|
||||
"GENERIC_USERINFO_ENDPOINT": "https://auth.example.com/userinfo",
|
||||
},
|
||||
200,
|
||||
"generic",
|
||||
[],
|
||||
),
|
||||
(
|
||||
{"GENERIC_CLIENT_ID": "test-generic-client-id"},
|
||||
503,
|
||||
"generic",
|
||||
[
|
||||
"GENERIC_CLIENT_SECRET",
|
||||
"GENERIC_AUTHORIZATION_ENDPOINT",
|
||||
"GENERIC_TOKEN_ENDPOINT",
|
||||
"GENERIC_USERINFO_ENDPOINT",
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_sso_readiness_generic_configurations(
|
||||
self, env_vars, expected_status, expected_provider, expected_missing_vars
|
||||
):
|
||||
"""Test Generic SSO readiness with both fully configured and missing variables"""
|
||||
from fastapi.testclient import TestClient
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
client = TestClient(app)
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=True):
|
||||
response = client.get("/sso/readiness")
|
||||
|
||||
assert response.status_code == expected_status
|
||||
|
||||
if expected_status == 200:
|
||||
data = response.json()
|
||||
assert data["sso_configured"] is True
|
||||
assert data["provider"] == expected_provider
|
||||
assert data["status"] == "healthy"
|
||||
assert "Generic SSO is properly configured" in data["message"]
|
||||
else:
|
||||
data = response.json()["detail"]
|
||||
assert data["sso_configured"] is True
|
||||
assert data["provider"] == expected_provider
|
||||
assert data["status"] == "unhealthy"
|
||||
assert set(data["missing_environment_variables"]) == set(
|
||||
expected_missing_vars
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user