mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
Refactor settings response models in proxy_setting_endpoints.py
- Renamed SSOSettingsResponse to inherit from a new base class SettingsResponse for better structure. - Introduced InternalUserSettingsResponse and DefaultTeamSettingsResponse models for internal user and default team settings. - Updated endpoint responses to use field_schema instead of schema for consistency. - Enhanced test cases to validate the new response structure and ensure proper functionality of SSO settings.
This commit is contained in:
Cole McIntosh
parent
3cc9460922
commit
3a946933ee
2 changed files with 121 additions and 22 deletions
|
|
@ -16,16 +16,31 @@ class IPAddress(BaseModel):
|
|||
ip: str
|
||||
|
||||
|
||||
class SSOSettingsResponse(BaseModel):
|
||||
"""Response model for SSO settings with values and schema information"""
|
||||
class SettingsResponse(BaseModel):
|
||||
"""Base response model for settings with values and schema information"""
|
||||
|
||||
values: Dict[str, Any]
|
||||
"""The current SSO configuration values"""
|
||||
"""The current configuration values"""
|
||||
|
||||
schema: Dict[str, Any]
|
||||
field_schema: Dict[str, Any]
|
||||
"""Schema information including descriptions and property types for UI display"""
|
||||
|
||||
|
||||
class SSOSettingsResponse(SettingsResponse):
|
||||
"""Response model for SSO settings"""
|
||||
pass
|
||||
|
||||
|
||||
class InternalUserSettingsResponse(SettingsResponse):
|
||||
"""Response model for internal user settings"""
|
||||
pass
|
||||
|
||||
|
||||
class DefaultTeamSettingsResponse(SettingsResponse):
|
||||
"""Response model for default team settings"""
|
||||
pass
|
||||
|
||||
|
||||
@router.get(
|
||||
"/get/allowed_ips",
|
||||
tags=["Budget & Spend Tracking"],
|
||||
|
|
@ -151,19 +166,19 @@ async def _get_settings_with_schema(
|
|||
# Add descriptions to the response
|
||||
result = {
|
||||
"values": settings_dict,
|
||||
"schema": {"description": schema.get("description", ""), "properties": {}},
|
||||
"field_schema": {"description": schema.get("description", ""), "properties": {}},
|
||||
}
|
||||
|
||||
# Add property descriptions
|
||||
for field_name, field_info in schema["properties"].items():
|
||||
result["schema"]["properties"][field_name] = {
|
||||
result["field_schema"]["properties"][field_name] = {
|
||||
"description": field_info.get("description", ""),
|
||||
"type": field_info.get("type", "string"),
|
||||
}
|
||||
|
||||
# Add nested object descriptions
|
||||
for def_name, def_schema in schema.get("definitions", {}).items():
|
||||
result["schema"][def_name] = {
|
||||
result["field_schema"][def_name] = {
|
||||
"description": def_schema.get("description", ""),
|
||||
"properties": {
|
||||
prop_name: {"description": prop_info.get("description", "")}
|
||||
|
|
@ -178,6 +193,7 @@ async def _get_settings_with_schema(
|
|||
"/get/internal_user_settings",
|
||||
tags=["SSO Settings"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=InternalUserSettingsResponse,
|
||||
)
|
||||
async def get_internal_user_settings():
|
||||
"""
|
||||
|
|
@ -200,6 +216,7 @@ async def get_internal_user_settings():
|
|||
"/get/default_team_settings",
|
||||
tags=["SSO Settings"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=DefaultTeamSettingsResponse,
|
||||
)
|
||||
async def get_default_team_settings():
|
||||
"""
|
||||
|
|
@ -345,12 +362,12 @@ async def get_sso_settings():
|
|||
# Add descriptions to the response
|
||||
result = {
|
||||
"values": sso_dict,
|
||||
"schema": {"description": schema.get("description", ""), "properties": {}},
|
||||
"field_schema": {"description": schema.get("description", ""), "properties": {}},
|
||||
}
|
||||
|
||||
# Add property descriptions
|
||||
for field_name, field_info in schema["properties"].items():
|
||||
result["schema"]["properties"][field_name] = {
|
||||
result["field_schema"]["properties"][field_name] = {
|
||||
"description": field_info.get("description", ""),
|
||||
"type": field_info.get("type", "string"),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ sys.path.insert(
|
|||
|
||||
from litellm.proxy._types import DefaultInternalUserParams, LitellmUserRoles
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams
|
||||
from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams, SSOConfig
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
|
|
@ -34,6 +34,16 @@ def mock_proxy_config(monkeypatch):
|
|||
"tpm_limit": 100,
|
||||
"rpm_limit": 10,
|
||||
},
|
||||
},
|
||||
"general_settings": {
|
||||
"proxy_admin_email": "admin@example.com"
|
||||
},
|
||||
"environment_variables": {
|
||||
"GOOGLE_CLIENT_ID": "test_google_client_id",
|
||||
"GOOGLE_CLIENT_SECRET": "test_google_client_secret",
|
||||
"MICROSOFT_CLIENT_ID": "test_microsoft_client_id",
|
||||
"MICROSOFT_CLIENT_SECRET": "test_microsoft_client_secret",
|
||||
"PROXY_BASE_URL": "https://example.com"
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -84,9 +94,9 @@ class TestProxySettingEndpoints:
|
|||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check structure of response
|
||||
# Check structure of response (updated to use field_schema)
|
||||
assert "values" in data
|
||||
assert "schema" in data
|
||||
assert "field_schema" in data
|
||||
|
||||
# Check values match our mock config
|
||||
values = data["values"]
|
||||
|
|
@ -98,10 +108,10 @@ class TestProxySettingEndpoints:
|
|||
assert values["budget_duration"] == mock_params["budget_duration"]
|
||||
assert values["models"] == mock_params["models"]
|
||||
|
||||
# Check schema contains descriptions
|
||||
assert "properties" in data["schema"]
|
||||
assert "user_role" in data["schema"]["properties"]
|
||||
assert "description" in data["schema"]["properties"]["user_role"]
|
||||
# Check field_schema contains descriptions (updated from schema to field_schema)
|
||||
assert "properties" in data["field_schema"]
|
||||
assert "user_role" in data["field_schema"]["properties"]
|
||||
assert "description" in data["field_schema"]["properties"]["user_role"]
|
||||
|
||||
def test_update_internal_user_settings(
|
||||
self, mock_proxy_config, mock_auth, monkeypatch
|
||||
|
|
@ -153,9 +163,9 @@ class TestProxySettingEndpoints:
|
|||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check structure of response
|
||||
# Check structure of response (updated to use field_schema)
|
||||
assert "values" in data
|
||||
assert "schema" in data
|
||||
assert "field_schema" in data
|
||||
|
||||
# Check values match our mock config
|
||||
values = data["values"]
|
||||
|
|
@ -168,10 +178,10 @@ class TestProxySettingEndpoints:
|
|||
assert values["tpm_limit"] == mock_params["tpm_limit"]
|
||||
assert values["rpm_limit"] == mock_params["rpm_limit"]
|
||||
|
||||
# Check schema contains descriptions
|
||||
assert "properties" in data["schema"]
|
||||
assert "models" in data["schema"]["properties"]
|
||||
assert "description" in data["schema"]["properties"]["models"]
|
||||
# Check field_schema contains descriptions (updated from schema to field_schema)
|
||||
assert "properties" in data["field_schema"]
|
||||
assert "models" in data["field_schema"]["properties"]
|
||||
assert "description" in data["field_schema"]["properties"]["models"]
|
||||
|
||||
def test_update_default_team_settings(
|
||||
self, mock_proxy_config, mock_auth, monkeypatch
|
||||
|
|
@ -218,3 +228,75 @@ class TestProxySettingEndpoints:
|
|||
|
||||
# Verify save_config was called exactly once
|
||||
assert mock_proxy_config["save_call_count"]() == 1
|
||||
|
||||
def test_get_sso_settings(self, mock_proxy_config, mock_auth):
|
||||
"""Test getting the SSO settings"""
|
||||
response = client.get("/get/sso_settings")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check structure of response
|
||||
assert "values" in data
|
||||
assert "field_schema" in data
|
||||
|
||||
# Check values contain SSO configuration
|
||||
values = data["values"]
|
||||
assert "google_client_id" in values
|
||||
assert "google_client_secret" in values
|
||||
assert "microsoft_client_id" in values
|
||||
assert "microsoft_client_secret" in values
|
||||
assert "proxy_base_url" in values
|
||||
assert "user_email" in values
|
||||
|
||||
# Verify values match our mock config
|
||||
assert values["google_client_id"] == "test_google_client_id"
|
||||
assert values["google_client_secret"] == "test_google_client_secret"
|
||||
assert values["microsoft_client_id"] == "test_microsoft_client_id"
|
||||
assert values["microsoft_client_secret"] == "test_microsoft_client_secret"
|
||||
assert values["proxy_base_url"] == "https://example.com"
|
||||
assert values["user_email"] == "admin@example.com"
|
||||
|
||||
# Check field_schema contains descriptions
|
||||
assert "properties" in data["field_schema"]
|
||||
assert "google_client_id" in data["field_schema"]["properties"]
|
||||
assert "description" in data["field_schema"]["properties"]["google_client_id"]
|
||||
|
||||
def test_update_sso_settings(self, mock_proxy_config, mock_auth):
|
||||
"""Test updating the SSO settings"""
|
||||
# New SSO settings to update
|
||||
new_sso_settings = {
|
||||
"google_client_id": "new_google_client_id",
|
||||
"google_client_secret": "new_google_client_secret",
|
||||
"microsoft_client_id": "new_microsoft_client_id",
|
||||
"microsoft_client_secret": "new_microsoft_client_secret",
|
||||
"proxy_base_url": "https://newexample.com",
|
||||
"user_email": "newadmin@example.com"
|
||||
}
|
||||
|
||||
response = client.patch("/update/sso_settings", json=new_sso_settings)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check response structure
|
||||
assert data["status"] == "success"
|
||||
assert "settings" in data
|
||||
|
||||
# Verify settings were updated
|
||||
settings = data["settings"]
|
||||
assert settings["google_client_id"] == new_sso_settings["google_client_id"]
|
||||
assert settings["google_client_secret"] == new_sso_settings["google_client_secret"]
|
||||
assert settings["microsoft_client_id"] == new_sso_settings["microsoft_client_id"]
|
||||
assert settings["microsoft_client_secret"] == new_sso_settings["microsoft_client_secret"]
|
||||
assert settings["proxy_base_url"] == new_sso_settings["proxy_base_url"]
|
||||
assert settings["user_email"] == new_sso_settings["user_email"]
|
||||
|
||||
# Verify the config was updated
|
||||
updated_config = mock_proxy_config["config"]
|
||||
assert updated_config["environment_variables"]["GOOGLE_CLIENT_ID"] == new_sso_settings["google_client_id"]
|
||||
assert updated_config["environment_variables"]["GOOGLE_CLIENT_SECRET"] == new_sso_settings["google_client_secret"]
|
||||
assert updated_config["general_settings"]["proxy_admin_email"] == new_sso_settings["user_email"]
|
||||
|
||||
# Verify save_config was called exactly once
|
||||
assert mock_proxy_config["save_call_count"]() == 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue