import asyncio import importlib import os import tempfile from unittest.mock import AsyncMock, MagicMock, Mock, patch from uuid import uuid4 import pytest import litellm from litellm.proxy._types import KeyManagementSystem from litellm.secret_managers.main import _should_read_secret_from_secret_manager, get_secret from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome class MockSecretClient: def get_secret(self, secret_name): return Mock(value="mocked_secret_value") @pytest.mark.asyncio async def test_azure_kms(): """ Basic asserts that the value from get secret is from Azure Key Vault when Key Management System is Azure Key Vault """ with patch("litellm.secret_manager_client", new=MockSecretClient()): litellm._key_management_system = KeyManagementSystem.AZURE_KEY_VAULT secret = get_secret(secret_name="ishaan-test-key") assert secret == "mocked_secret_value" @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): """ This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. """ importlib.reload(litellm) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop) yield loop.close() asyncio.set_event_loop(None) def redact_oidc_signature(secret_val: str) -> list[str]: return secret_val.split(".")[:-1] + ["SIGNATURE_REMOVED"] def test_oidc_env_variable(): # Create a unique environment variable name env_var_name = "OIDC_TEST_PATH_" + uuid4().hex os.environ[env_var_name] = "secret-" + uuid4().hex secret_val = get_secret(f"oidc/env/{env_var_name}") print(f"secret_val: {redact_oidc_signature(secret_val)}") assert secret_val == os.environ[env_var_name] # now unset the environment variable del os.environ[env_var_name] def test_oidc_file(monkeypatch): with tempfile.TemporaryDirectory() as temp_dir: monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", temp_dir) temp_file_path = os.path.join(temp_dir, "token.txt") secret_value = "secret-" + uuid4().hex with open(temp_file_path, "w") as temp_file: temp_file.write(secret_value) secret_val = get_secret(f"oidc/file/{temp_file_path}") print(f"secret_val: {redact_oidc_signature(secret_val)}") assert secret_val == secret_value def test_oidc_env_path(): # Create a temporary file with tempfile.NamedTemporaryFile(mode="w+") as temp_file: secret_value = "secret-" + uuid4().hex temp_file.write(secret_value) temp_file.flush() temp_file_path = temp_file.name # Create a unique environment variable name env_var_name = "OIDC_TEST_PATH_" + uuid4().hex # Set the environment variable to the temporary file path os.environ[env_var_name] = temp_file_path # Test getting the secret using the environment variable secret_val = get_secret(f"oidc/env_path/{env_var_name}") print(f"secret_val: {redact_oidc_signature(secret_val)}") assert secret_val == secret_value del os.environ[env_var_name] def test_should_read_secret_from_secret_manager(): """ Test that _should_read_secret_from_secret_manager returns correct values based on access mode """ from litellm.types.secret_managers.main import KeyManagementSettings # Test when secret manager client is None litellm.secret_manager_client = None litellm._key_management_settings = KeyManagementSettings() assert _should_read_secret_from_secret_manager() is False # Test with secret manager client and read_only access litellm.secret_manager_client = "dummy_client" litellm._key_management_settings = KeyManagementSettings(access_mode="read_only") assert _should_read_secret_from_secret_manager() is True # Test with secret manager client and read_and_write access litellm._key_management_settings = KeyManagementSettings( access_mode="read_and_write" ) assert _should_read_secret_from_secret_manager() is True # Test with secret manager client and write_only access litellm._key_management_settings = KeyManagementSettings(access_mode="write_only") assert _should_read_secret_from_secret_manager() is False # Reset global variables litellm.secret_manager_client = None litellm._key_management_settings = KeyManagementSettings() def test_get_secret_with_access_mode(): """ Test that get_secret respects access mode settings """ from litellm.types.secret_managers.main import KeyManagementSettings # Set up test environment test_secret_name = "TEST_SECRET_KEY" test_secret_value = "test_secret_value" os.environ[test_secret_name] = test_secret_value # Test with write_only access (should read from os.environ) litellm.secret_manager_client = "dummy_client" litellm._key_management_settings = KeyManagementSettings(access_mode="write_only") assert get_secret(test_secret_name) == test_secret_value # Test with no KeyManagementSettings but secret_manager_client set litellm.secret_manager_client = "dummy_client" litellm._key_management_settings = KeyManagementSettings() assert _should_read_secret_from_secret_manager() is True # Test with read_only access litellm._key_management_settings = KeyManagementSettings(access_mode="read_only") assert _should_read_secret_from_secret_manager() is True # Test with read_and_write access litellm._key_management_settings = KeyManagementSettings( access_mode="read_and_write" ) assert _should_read_secret_from_secret_manager() is True # Reset global variables litellm.secret_manager_client = None litellm._key_management_settings = KeyManagementSettings() del os.environ[test_secret_name] def test_key_management_settings_defaults(): """ Test that KeyManagementSettings initializes with correct default values. """ from litellm.types.secret_managers.main import KeyManagementSettings settings = KeyManagementSettings() assert settings.store_virtual_keys is False assert settings.prefix_for_stored_virtual_keys == "litellm/" assert settings.access_mode == "read_only" assert settings.description is None assert settings.tags is None assert settings.primary_secret_name is None def test_key_management_settings_custom_values(): """ Test that KeyManagementSettings correctly stores custom description and tags. """ from litellm.types.secret_managers.main import KeyManagementSettings custom_tags = {"Environment": "Dev", "Team": "Intelligence"} custom_description = "LiteLLM-managed API key for development" settings = KeyManagementSettings( store_virtual_keys=True, prefix_for_stored_virtual_keys="litellm/custom/", access_mode="read_and_write", primary_secret_name="primary/litellm/keys", description=custom_description, tags=custom_tags, ) assert settings.store_virtual_keys is True assert settings.prefix_for_stored_virtual_keys == "litellm/custom/" assert settings.access_mode == "read_and_write" assert settings.primary_secret_name == "primary/litellm/keys" assert settings.description == custom_description assert settings.tags == custom_tags @pytest.mark.asyncio async def test_async_write_secret_receives_description_and_tags(monkeypatch): """ Test that AWSSecretsManagerV2.async_write_secret receives description and tags when KeyManagementSettings is set. """ from litellm import litellm from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 from litellm.types.secret_managers.main import KeyManagementSettings # Mock out AWS network calls mock_async_write = AsyncMock(return_value={"Name": "litellm/test_secret"}) monkeypatch.setattr(AWSSecretsManagerV2, "async_write_secret", mock_async_write) # Setup settings litellm._key_management_settings = KeyManagementSettings( store_virtual_keys=True, description="LiteLLM Unit Test Secret", tags={"Owner": "UnitTest", "Purpose": "Validation"}, ) # Instantiate fake client litellm.secret_manager_client = AWSSecretsManagerV2() # Call the helper method that stores a virtual key from litellm.proxy.hooks.key_management_event_hooks import ( KeyManagementEventHooks, ) await KeyManagementEventHooks._store_virtual_key_in_secret_manager( secret_name="test_secret", secret_token="test_value" ) # Verify async_write_secret was called with correct metadata mock_async_write.assert_called_once() args, kwargs = mock_async_write.call_args assert kwargs["secret_name"].endswith("test_secret") assert kwargs["secret_value"] == "test_value" assert kwargs["description"] == "LiteLLM Unit Test Secret" assert kwargs["tags"] == {"Owner": "UnitTest", "Purpose": "Validation"} def test_key_management_settings_serialization_roundtrip(): """ Test that KeyManagementSettings serializes and deserializes consistently (Pydantic behavior). """ from litellm.types.secret_managers.main import KeyManagementSettings original = KeyManagementSettings( store_virtual_keys=True, prefix_for_stored_virtual_keys="litellm/dev/", access_mode="read_and_write", description="Roundtrip test", tags={"Env": "QA"}, ) as_dict = original.model_dump() reloaded = KeyManagementSettings(**as_dict) assert reloaded.store_virtual_keys is True assert reloaded.prefix_for_stored_virtual_keys == "litellm/dev/" assert reloaded.access_mode == "read_and_write" assert reloaded.description == "Roundtrip test" assert reloaded.tags == {"Env": "QA"}