tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 08:25:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Infra] Litellm Backend SSO Changes (#16029)

Browse Source 
* SSO Backend changes

* Encrypt and Decrypt, load into os env

* Linting and addressing comments
This commit is contained in:
yuneng-jiang
2025-10-30 14:32:08 -07:00
committed by GitHub
parent eed3ad0bdb
commit 720ba865fb
7 changed files with 617 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+8
View File
@@ -580,4 +580,12 @@ model LiteLLM_SearchToolsTable {
search_tool_info Json?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
// SSO configuration table
model LiteLLM_SSOConfig {
id String @id @default("sso_config")
sso_settings Json
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
litellm/proxy/proxy_server.py
+23
View File
@@ -1071,6 +1071,7 @@ celery_fn = None # Redis Queue for handling requests
scheduler = None
last_model_cost_map_reload = None
### DB WRITER ###
db_writer_client: Optional[AsyncHTTPHandler] = None
### logger ###
@@ -3335,6 +3336,28 @@ class ProxyConfig:
if self._should_load_db_object(object_type="model_cost_map"):
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
if self._should_load_db_object(object_type="sso_settings"):
await self._init_sso_settings_in_db(prisma_client=prisma_client)
async def _init_sso_settings_in_db(self, prisma_client: PrismaClient):
"""
Initialize SSO settings from database into the router on startup.
"""
try:
sso_settings = await prisma_client.db.litellm_ssoconfig.find_unique(
where={"id": "sso_config"}
)
if sso_settings is not None:
# Capitalize all keys in sso_settings dictionary
uppercase_sso_settings = {key.upper(): value for key, value in sso_settings.sso_settings.items()}
self._decrypt_and_set_db_env_variables(environment_variables=uppercase_sso_settings)
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_sso_settings_in_db - {}".format(
str(e)
)
)
async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient):
"""
litellm/proxy/schema.prisma
+8
View File
@@ -580,4 +580,12 @@ model LiteLLM_SearchToolsTable {
search_tool_info Json?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
// SSO configuration table
model LiteLLM_SSOConfig {
id String @id @default("sso_config")
sso_settings Json
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
+70 -58
View File
@@ -390,40 +390,47 @@ async def update_default_team_settings(settings: DefaultTeamSSOParams):
)
async def get_sso_settings():
"""
Get all SSO configuration settings from the environment variables.
Get all SSO configuration settings from the dedicated SSO table.
Returns a structured object with values and descriptions for UI display.
"""
import os
from litellm.proxy.proxy_server import proxy_config
from litellm.proxy.proxy_server import prisma_client, proxy_config
# Load existing config to get both environment variables and general settings
config = await proxy_config.get_config()
general_settings = config.get("general_settings", {}) or {}
environment_variables = config.get("environment_variables", {}) or {}
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": "Database not connected. Please connect a database."},
)
# Get user_email from general_settings
proxy_admin_email = general_settings.get("proxy_admin_email", None)
# Get SSO config from dedicated table
sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique(
where={"id": "sso_config"}
)
# Helper function to get env var value (first from config, then from environment)
def get_env_value(env_var_name: str):
return environment_variables.get(env_var_name) or os.getenv(env_var_name)
# Initialize with defaults
sso_settings_dict = {}
if sso_db_record and sso_db_record.sso_settings:
# Load settings from database
sso_settings_dict = dict(sso_db_record.sso_settings)
decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables(environment_variables=sso_settings_dict)
# Get current environment variables for SSO
# Build SSO config with database values or environment fallback
sso_config = SSOConfig(
google_client_id=get_env_value("GOOGLE_CLIENT_ID"),
google_client_secret=get_env_value("GOOGLE_CLIENT_SECRET"),
microsoft_client_id=get_env_value("MICROSOFT_CLIENT_ID"),
microsoft_client_secret=get_env_value("MICROSOFT_CLIENT_SECRET"),
microsoft_tenant=get_env_value("MICROSOFT_TENANT"),
generic_client_id=get_env_value("GENERIC_CLIENT_ID"),
generic_client_secret=get_env_value("GENERIC_CLIENT_SECRET"),
generic_authorization_endpoint=get_env_value("GENERIC_AUTHORIZATION_ENDPOINT"),
generic_token_endpoint=get_env_value("GENERIC_TOKEN_ENDPOINT"),
generic_userinfo_endpoint=get_env_value("GENERIC_USERINFO_ENDPOINT"),
proxy_base_url=get_env_value("PROXY_BASE_URL"),
user_email=proxy_admin_email, # Get from config instead of environment
ui_access_mode=general_settings.get("ui_access_mode", None),
google_client_id=decrypted_sso_settings_dict.get("google_client_id", None),
google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None),
microsoft_client_id=decrypted_sso_settings_dict.get("microsoft_client_id", None),
microsoft_client_secret=decrypted_sso_settings_dict.get("microsoft_client_secret", None),
microsoft_tenant=decrypted_sso_settings_dict.get("microsoft_tenant", None),
generic_client_id=decrypted_sso_settings_dict.get("generic_client_id", None),
generic_client_secret=decrypted_sso_settings_dict.get("generic_client_secret", None),
generic_authorization_endpoint=decrypted_sso_settings_dict.get("generic_authorization_endpoint", None),
generic_token_endpoint=decrypted_sso_settings_dict.get("generic_token_endpoint", None),
generic_userinfo_endpoint=decrypted_sso_settings_dict.get("generic_userinfo_endpoint", None),
proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None),
user_email=decrypted_sso_settings_dict.get("user_email"),
ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"),
)
# Get the schema for UI display
@@ -460,11 +467,26 @@ async def get_sso_settings():
)
async def update_sso_settings(sso_config: SSOConfig):
"""
Update SSO configuration by saving to both environment variables and config file.
Update SSO configuration by saving to the dedicated SSO table.
"""
import os
import json
from litellm.proxy.proxy_server import proxy_config
from litellm.proxy.proxy_server import prisma_client, store_model_in_db, proxy_config
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": "Database not connected. Please connect a database."},
)
if store_model_in_db is not True:
raise HTTPException(
status_code=500,
detail={
"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."
},
)
# Update environment variables
env_var_mapping = {
@@ -495,39 +517,29 @@ async def update_sso_settings(sso_config: SSOConfig):
# Update environment variables in config and in memory
sso_data = sso_config.model_dump()
for field_name, value in sso_data.items():
if field_name == "user_email":
if value:
# Store user_email in general_settings instead of environment variables
config["general_settings"]["proxy_admin_email"] = value
else:
# Clear user_email if null/empty
config["general_settings"].pop("proxy_admin_email", None)
elif field_name == "ui_access_mode":
if value:
config["general_settings"]["ui_access_mode"] = value
else:
# Clear ui_access_mode if null/empty
config["general_settings"].pop("ui_access_mode", None)
elif field_name in env_var_mapping and value:
if field_name in env_var_mapping:
env_var_name = env_var_mapping[field_name]
# Update in config
config["environment_variables"][env_var_name] = value
# Update in runtime environment
os.environ[env_var_name] = value
elif field_name in env_var_mapping:
# Clear environment variable if value is null/empty
env_var_name = env_var_mapping[field_name]
config["environment_variables"].pop(env_var_name, None)
os.environ.pop(env_var_name, None)
if value:
os.environ[env_var_name] = value
else:
# Clear environment variable if value is null/empty
os.environ.pop(env_var_name, None)
stored_config = config
if len(config["environment_variables"]) > 0:
encrypted_sso_data = proxy_config._encrypt_env_variables(environment_variables=sso_data)
stored_config["environment_variables"] = proxy_config._encrypt_env_variables(
environment_variables=config["environment_variables"]
)
# Save the updated config
await proxy_config.save_config(new_config=stored_config)
# Save to dedicated SSO table
await prisma_client.db.litellm_ssoconfig.upsert(
where={"id": "sso_config"},
data={
"create": {
"id": "sso_config",
"sso_settings": json.dumps(encrypted_sso_data),
},
"update": {
"sso_settings": json.dumps(encrypted_sso_data),
},
},
)
return {
"message": "SSO settings updated successfully",
schema.prisma
+8
View File
@@ -580,4 +580,12 @@ model LiteLLM_SearchToolsTable {
search_tool_info Json?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
// SSO configuration table
model LiteLLM_SSOConfig {
id String @id @default("sso_config")
sso_settings Json
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
tests/test_litellm/proxy/test_proxy_server.py
+155
View File
@@ -2288,3 +2288,158 @@ async def test_tag_cache_update_multiple_tags():
assert tag_updates["tag:tag2"]["spend"] == 25.0
@pytest.mark.asyncio
async def test_init_sso_settings_in_db():
"""
Test that _init_sso_settings_in_db properly loads SSO settings from database,
uppercases keys, and calls _decrypt_and_set_db_env_variables.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
# Test Case 1: SSO settings exist in database
mock_sso_config = MagicMock()
mock_sso_config.sso_settings = {
"google_client_id": "test-client-id",
"google_client_secret": "test-client-secret",
"microsoft_client_id": "ms-client-id",
"microsoft_client_secret": "ms-client-secret",
}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(
return_value=mock_sso_config
)
# Mock _decrypt_and_set_db_env_variables
with patch.object(
proxy_config, "_decrypt_and_set_db_env_variables"
) as mock_decrypt_and_set:
await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client)
# Verify find_unique was called with correct parameters
mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(
where={"id": "sso_config"}
)
# Verify _decrypt_and_set_db_env_variables was called with uppercased keys
mock_decrypt_and_set.assert_called_once()
call_args = mock_decrypt_and_set.call_args
uppercased_settings = call_args.kwargs["environment_variables"]
# Verify all keys are uppercased
assert "GOOGLE_CLIENT_ID" in uppercased_settings
assert "GOOGLE_CLIENT_SECRET" in uppercased_settings
assert "MICROSOFT_CLIENT_ID" in uppercased_settings
assert "MICROSOFT_CLIENT_SECRET" in uppercased_settings
# Verify values are preserved
assert uppercased_settings["GOOGLE_CLIENT_ID"] == "test-client-id"
assert uppercased_settings["GOOGLE_CLIENT_SECRET"] == "test-client-secret"
assert uppercased_settings["MICROSOFT_CLIENT_ID"] == "ms-client-id"
assert uppercased_settings["MICROSOFT_CLIENT_SECRET"] == "ms-client-secret"
# Verify original lowercase keys are not present
assert "google_client_id" not in uppercased_settings
assert "microsoft_client_id" not in uppercased_settings
@pytest.mark.asyncio
async def test_init_sso_settings_in_db_no_settings():
"""
Test that _init_sso_settings_in_db handles the case when no SSO settings exist in database.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
# Mock prisma client to return None (no SSO settings)
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None)
# Mock _decrypt_and_set_db_env_variables
with patch.object(
proxy_config, "_decrypt_and_set_db_env_variables"
) as mock_decrypt_and_set:
await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client)
# Verify find_unique was called
mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(
where={"id": "sso_config"}
)
# Verify _decrypt_and_set_db_env_variables was NOT called when no settings exist
mock_decrypt_and_set.assert_not_called()
@pytest.mark.asyncio
async def test_init_sso_settings_in_db_error_handling():
"""
Test that _init_sso_settings_in_db handles database errors gracefully.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
# Mock prisma client to raise an exception
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(
side_effect=Exception("Database connection error")
)
# The method should not raise an exception, it should log it instead
try:
await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client)
# If we get here, the exception was handled properly
assert True
except Exception as e:
# The exception should be caught and logged, not propagated
pytest.fail(f"Exception should have been caught and logged, but was raised: {e}")
@pytest.mark.asyncio
async def test_init_sso_settings_in_db_empty_settings():
"""
Test that _init_sso_settings_in_db handles empty SSO settings dictionary.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
# Mock SSO config with empty settings dictionary
mock_sso_config = MagicMock()
mock_sso_config.sso_settings = {}
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(
return_value=mock_sso_config
)
# Mock _decrypt_and_set_db_env_variables
with patch.object(
proxy_config, "_decrypt_and_set_db_env_variables"
) as mock_decrypt_and_set:
await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client)
# Verify find_unique was called
mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(
where={"id": "sso_config"}
)
# Verify _decrypt_and_set_db_env_variables was called with empty dict
mock_decrypt_and_set.assert_called_once()
call_args = mock_decrypt_and_set.call_args
uppercased_settings = call_args.kwargs["environment_variables"]
# Verify empty dictionary
assert uppercased_settings == {}
tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
+345 -99
View File
@@ -234,8 +234,31 @@ 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"""
def test_get_sso_settings(self, mock_proxy_config, mock_auth, monkeypatch):
"""Test getting the SSO settings from the dedicated database table"""
from unittest.mock import AsyncMock, MagicMock
# Mock the prisma client with database record
# Note: Prisma returns Json fields as dicts (auto-parsed)
mock_prisma = MagicMock()
mock_db_record = MagicMock()
mock_db_record.sso_settings = {
"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",
"user_email": "admin@example.com",
}
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock decryption to return the values as-is (simulating decryption)
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(
proxy_config, "_decrypt_and_set_db_env_variables", lambda environment_variables: environment_variables
)
response = client.get("/get/sso_settings")
assert response.status_code == 200
@@ -266,10 +289,29 @@ class TestProxySettingEndpoints:
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"]
# Verify find_unique was called with correct parameters
mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once()
call_args = mock_prisma.db.litellm_ssoconfig.find_unique.call_args
assert call_args.kwargs["where"]["id"] == "sso_config"
def test_update_sso_settings(self, mock_proxy_config, mock_auth, monkeypatch):
"""Test updating the SSO settings to the dedicated database table"""
import json
from unittest.mock import AsyncMock, MagicMock
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
"""Test updating the SSO settings"""
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
# Mock the prisma client
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock encryption to return values as-is
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables)
# New SSO settings to update
new_sso_settings = {
"google_client_id": "new_google_client_id",
@@ -305,34 +347,44 @@ class TestProxySettingEndpoints:
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
# Verify upsert was called with correct parameters
assert mock_prisma.db.litellm_ssoconfig.upsert.called
call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args
# Verify the upsert is using the correct ID
assert call_args.kwargs["where"]["id"] == "sso_config"
# Verify the data structure for create and update
create_data = call_args.kwargs["data"]["create"]
update_data = call_args.kwargs["data"]["update"]
assert create_data["id"] == "sso_config"
assert "sso_settings" in create_data
assert "sso_settings" in update_data
# Verify the data is stored as JSON string (as per implementation)
# The encryption mock returns data as-is, so we verify structure
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["google_client_id"] == "new_google_client_id"
def test_update_sso_settings_with_null_values_clears_env_vars(
self, mock_proxy_config, mock_auth, monkeypatch
):
"""Test that updating SSO settings with null values clears environment variables"""
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
"""Test that updating SSO settings with null values clears environment variables and updates database"""
import json
from unittest.mock import AsyncMock, MagicMock
# First, verify we have existing environment variables
initial_config = mock_proxy_config["config"]
assert "GOOGLE_CLIENT_ID" in initial_config["environment_variables"]
assert "MICROSOFT_CLIENT_ID" in initial_config["environment_variables"]
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
# Mock the prisma client
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock encryption to return values as-is
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables)
# Set some initial environment variables for runtime testing
monkeypatch.setenv("GOOGLE_CLIENT_ID", "test_existing_google_id")
@@ -361,32 +413,38 @@ class TestProxySettingEndpoints:
data = response.json()
assert data["status"] == "success"
# Verify that environment variables were cleared from config
updated_config = mock_proxy_config["config"]
# These should be removed from environment_variables
assert "GOOGLE_CLIENT_ID" not in updated_config["environment_variables"]
assert "GOOGLE_CLIENT_SECRET" not in updated_config["environment_variables"]
assert "MICROSOFT_CLIENT_ID" not in updated_config["environment_variables"]
assert "MICROSOFT_CLIENT_SECRET" not in updated_config["environment_variables"]
assert "MICROSOFT_TENANT" not in updated_config["environment_variables"]
assert "PROXY_BASE_URL" not in updated_config["environment_variables"]
# Verify that runtime environment variables were cleared
assert "GOOGLE_CLIENT_ID" not in os.environ
assert "MICROSOFT_CLIENT_ID" not in os.environ
# Verify user_email was cleared from general_settings
assert updated_config["general_settings"].get("proxy_admin_email") is None
# Verify save_config was called
assert mock_proxy_config["save_call_count"]() == 1
# Verify upsert was called with correct parameters
assert mock_prisma.db.litellm_ssoconfig.upsert.called
call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args
# Verify null values are stored in database
create_data = call_args.kwargs["data"]["create"]
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["google_client_id"] is None
assert create_sso_settings["microsoft_client_id"] is None
def test_update_sso_settings_with_empty_strings_clears_env_vars(
self, mock_proxy_config, mock_auth, monkeypatch
):
"""Test that updating SSO settings with empty strings also clears environment variables"""
"""Test that updating SSO settings with empty strings also clears environment variables and updates database"""
import json
from unittest.mock import AsyncMock, MagicMock
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
# Mock the prisma client
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock encryption to return values as-is
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables)
# Set some initial environment variables for runtime testing
monkeypatch.setenv("GOOGLE_CLIENT_ID", "test_existing_google_id")
@@ -407,28 +465,38 @@ class TestProxySettingEndpoints:
data = response.json()
assert data["status"] == "success"
# Verify that environment variables with empty strings were cleared from config
updated_config = mock_proxy_config["config"]
assert "GOOGLE_CLIENT_ID" not in updated_config["environment_variables"]
assert "GOOGLE_CLIENT_SECRET" not in updated_config["environment_variables"]
assert "MICROSOFT_CLIENT_SECRET" not in updated_config["environment_variables"]
assert "PROXY_BASE_URL" not in updated_config["environment_variables"]
# Verify that runtime environment variables were cleared
assert "GOOGLE_CLIENT_ID" not in os.environ
assert "MICROSOFT_CLIENT_SECRET" not in os.environ
# Verify user_email was cleared from general_settings
assert updated_config["general_settings"].get("proxy_admin_email") is None
# Verify save_config was called
assert mock_proxy_config["save_call_count"]() == 1
# Verify upsert was called with correct parameters
assert mock_prisma.db.litellm_ssoconfig.upsert.called
call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args
# Verify empty strings are stored in database
create_data = call_args.kwargs["data"]["create"]
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["google_client_id"] == ""
assert create_sso_settings["microsoft_client_secret"] == ""
def test_update_sso_settings_mixed_null_and_valid_values(
self, mock_proxy_config, mock_auth, monkeypatch
):
"""Test updating SSO settings with mix of null and valid values"""
"""Test updating SSO settings with mix of null and valid values - verifies both env vars and database"""
import json
from unittest.mock import AsyncMock, MagicMock
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
# Mock the prisma client
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock encryption to return values as-is
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables)
# Set some initial environment variables
monkeypatch.setenv("GOOGLE_CLIENT_ID", "old_google_id")
@@ -451,44 +519,43 @@ class TestProxySettingEndpoints:
data = response.json()
assert data["status"] == "success"
# Verify the config was updated correctly
updated_config = mock_proxy_config["config"]
# Valid values should be set
assert (
updated_config["environment_variables"]["GOOGLE_CLIENT_ID"]
!= "new_google_client_id"
) # Encrypted
assert (
updated_config["environment_variables"]["MICROSOFT_CLIENT_SECRET"]
!= "new_microsoft_secret"
) # Encrypted
assert (
updated_config["environment_variables"]["PROXY_BASE_URL"]
!= "https://newproxy.com"
) # Encrypted
# Null values should be cleared
assert "GOOGLE_CLIENT_SECRET" not in updated_config["environment_variables"]
assert "MICROSOFT_CLIENT_ID" not in updated_config["environment_variables"]
# Verify runtime environment variables
assert os.environ.get("GOOGLE_CLIENT_ID") == "new_google_client_id"
assert os.environ.get("MICROSOFT_CLIENT_SECRET") == "new_microsoft_secret"
assert "GOOGLE_CLIENT_SECRET" not in os.environ
assert "MICROSOFT_CLIENT_ID" not in os.environ
# Verify user_email was cleared from general_settings
assert updated_config["general_settings"].get("proxy_admin_email") is None
# Verify save_config was called
assert mock_proxy_config["save_call_count"]() == 1
# Verify upsert was called with correct parameters
assert mock_prisma.db.litellm_ssoconfig.upsert.called
call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args
# Verify the mixed values are stored correctly in database
create_data = call_args.kwargs["data"]["create"]
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["google_client_id"] == "new_google_client_id"
assert create_sso_settings["google_client_secret"] is None
assert create_sso_settings["microsoft_client_id"] is None
assert create_sso_settings["microsoft_client_secret"] == "new_microsoft_secret"
assert create_sso_settings["proxy_base_url"] == "https://newproxy.com"
def test_update_sso_settings_ui_access_mode_handling(
self, mock_proxy_config, mock_auth, monkeypatch
):
"""Test that ui_access_mode is handled correctly in general_settings"""
"""Test that ui_access_mode is handled correctly and stored in database"""
import json
from unittest.mock import AsyncMock, MagicMock
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
# Mock the prisma client
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock encryption to return values as-is
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables)
# Test setting ui_access_mode
sso_settings_with_ui_mode = {
@@ -502,15 +569,13 @@ class TestProxySettingEndpoints:
data = response.json()
assert data["status"] == "success"
# Verify ui_access_mode was set in general_settings (not environment_variables)
updated_config = mock_proxy_config["config"]
assert updated_config["general_settings"]["ui_access_mode"] == "admin_only"
assert (
updated_config["general_settings"]["proxy_admin_email"] == "admin@test.com"
)
# Verify ui_access_mode is NOT in environment_variables
assert "ui_access_mode" not in updated_config["environment_variables"]
# Verify upsert was called with correct data
assert mock_prisma.db.litellm_ssoconfig.upsert.called
call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args
create_data = call_args.kwargs["data"]["create"]
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["ui_access_mode"] == "admin_only"
assert create_sso_settings["user_email"] == "admin@test.com"
# Test clearing ui_access_mode
clear_ui_mode = {"ui_access_mode": None, "user_email": None}
@@ -519,13 +584,13 @@ class TestProxySettingEndpoints:
assert response.status_code == 200
# Verify ui_access_mode and user_email were cleared
updated_config = mock_proxy_config["config"]
assert updated_config["general_settings"].get("ui_access_mode") is None
assert updated_config["general_settings"].get("proxy_admin_email") is None
# Verify save_config was called twice (once for each update)
assert mock_proxy_config["save_call_count"]() == 2
# Verify upsert was called again with null values
assert mock_prisma.db.litellm_ssoconfig.upsert.call_count == 2
call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args
create_data = call_args.kwargs["data"]["create"]
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["ui_access_mode"] is None
assert create_sso_settings["user_email"] is None
def test_get_ui_theme_settings(self, mock_proxy_config):
"""Test getting UI theme settings without authentication"""
@@ -556,3 +621,184 @@ class TestProxySettingEndpoints:
updated_config = mock_proxy_config["config"]
assert "UI_LOGO_PATH" in updated_config["environment_variables"]
assert mock_proxy_config["save_call_count"]() == 1
def test_get_sso_settings_from_database(self, mock_proxy_config, mock_auth, monkeypatch):
"""Test getting SSO settings from the dedicated database table"""
import json
from unittest.mock import AsyncMock, MagicMock
# Mock the prisma client
mock_prisma = MagicMock()
mock_db_record = MagicMock()
# Simulate encrypted data from database
mock_sso_settings = {
"google_client_id": "encrypted_google_id",
"google_client_secret": "encrypted_google_secret",
"microsoft_client_id": "encrypted_microsoft_id",
"proxy_base_url": "encrypted_proxy_url",
}
mock_db_record.sso_settings = mock_sso_settings
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock the decryption method to return decrypted values
def mock_decrypt_and_set(environment_variables):
return {
"google_client_id": "decrypted_google_id",
"google_client_secret": "decrypted_google_secret",
"microsoft_client_id": "decrypted_microsoft_id",
"proxy_base_url": "https://decrypted.example.com",
}
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(
proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt_and_set
)
response = client.get("/get/sso_settings")
assert response.status_code == 200
data = response.json()
# Verify structure
assert "values" in data
assert "field_schema" in data
# Verify decrypted values are returned
values = data["values"]
assert values["google_client_id"] == "decrypted_google_id"
assert values["google_client_secret"] == "decrypted_google_secret"
assert values["microsoft_client_id"] == "decrypted_microsoft_id"
assert values["proxy_base_url"] == "https://decrypted.example.com"
def test_update_sso_settings_to_database(self, mock_proxy_config, mock_auth, monkeypatch):
"""Test updating SSO settings saves to the dedicated database table"""
import json
from unittest.mock import AsyncMock, MagicMock
monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key")
# Mock the prisma client
mock_prisma = MagicMock()
upsert_mock = AsyncMock()
mock_prisma.db.litellm_ssoconfig.upsert = upsert_mock
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
# Track what was encrypted
encrypted_data = {}
def mock_encrypt(environment_variables):
# Simulate encryption by adding prefix
encrypted = {
k: f"encrypted_{v}" if v else v
for k, v in environment_variables.items()
}
encrypted_data.update(encrypted)
return encrypted
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(proxy_config, "_encrypt_env_variables", mock_encrypt)
# New SSO settings to save
new_sso_settings = {
"google_client_id": "new_google_id",
"google_client_secret": "new_google_secret",
"microsoft_client_id": "new_microsoft_id",
"proxy_base_url": "https://new.example.com",
}
response = client.patch("/update/sso_settings", json=new_sso_settings)
assert response.status_code == 200
data = response.json()
assert data["status"] == "success"
assert data["settings"]["google_client_id"] == "new_google_id"
# Verify upsert was called
assert upsert_mock.called
call_args = upsert_mock.call_args
# Verify it's using the correct ID
assert call_args.kwargs["where"]["id"] == "sso_config"
# Verify encrypted data was saved
create_data = call_args.kwargs["data"]["create"]
update_data = call_args.kwargs["data"]["update"]
assert create_data["id"] == "sso_config"
# The sso_settings should be JSON string of encrypted data
assert "sso_settings" in create_data
assert "sso_settings" in update_data
# Verify the encrypted data is correctly stored
create_sso_settings = json.loads(create_data["sso_settings"])
assert create_sso_settings["google_client_id"] == "encrypted_new_google_id"
assert create_sso_settings["google_client_secret"] == "encrypted_new_google_secret"
assert create_sso_settings["proxy_base_url"] == "encrypted_https://new.example.com"
def test_get_sso_settings_empty_database(self, mock_proxy_config, mock_auth, monkeypatch):
"""Test getting SSO settings when database table is empty"""
from unittest.mock import AsyncMock, MagicMock
# Mock the prisma client to return None (no record found)
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
# Mock the decryption method
def mock_decrypt_and_set(environment_variables):
# Should receive empty dict
return environment_variables
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(
proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt_and_set
)
response = client.get("/get/sso_settings")
assert response.status_code == 200
data = response.json()
# Verify structure is still correct with empty values
assert "values" in data
assert "field_schema" in data
# All values should be None
values = data["values"]
assert values.get("google_client_id") is None
assert values.get("google_client_secret") is None
assert values.get("microsoft_client_id") is None
def test_update_sso_settings_no_database_connection(self, mock_proxy_config, mock_auth, monkeypatch):
"""Test updating SSO settings when database is not connected"""
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
new_sso_settings = {
"google_client_id": "new_google_id",
}
response = client.patch("/update/sso_settings", json=new_sso_settings)
assert response.status_code == 500
data = response.json()
assert "error" in data["detail"]
assert "Database not connected" in data["detail"]["error"]
def test_get_sso_settings_no_database_connection(self, mock_proxy_config, mock_auth, monkeypatch):
"""Test getting SSO settings when database is not connected"""
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
response = client.get("/get/sso_settings")
assert response.status_code == 500
data = response.json()
assert "error" in data["detail"]
assert "Database not connected" in data["detail"]["error"]