(Feat) Hashicorp Secret Manager - Allow storing virtual keys in secret manager (#7549)
All checks were successful
Read Version from pyproject.toml / read-version (push) Successful in 13s

* use a base abstract class

* async_write_secret for hcorp

* fix hcorp

* async_write_secret for hashicopr secret manager

* store virtual keys in hcorp

* add delete secret

* test_hashicorp_secret_manager_write_secret

* test_hashicorp_secret_manager_delete_secret

* docs Supported Secret Managers

* docs storing keys in hcorp

* docs hcorp

* docs secret managers

* test_key_generate_with_secret_manager_call

* fix unused imports
This commit is contained in:
Ishaan Jaff 2025-01-04 11:35:59 -08:00 committed by GitHub
parent 7f7222ce30
commit 46d9d29bff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 458 additions and 119 deletions

View file

@ -21,7 +21,7 @@ This covers:
- ✅ [Audit Logs with retention policy](./proxy/enterprise#audit-logs)
- ✅ [JWT-Auth](../docs/proxy/token_auth.md)
- ✅ [Control available public, private routes (Restrict certain endpoints on proxy)](./proxy/enterprise#control-available-public-private-routes)
- ✅ [**Secret Managers** AWS Key Manager, Google Secret Manager, Azure Key](./secret)
- ✅ [**Secret Managers** AWS Key Manager, Google Secret Manager, Azure Key, Hashicorp Vault](./secret)
- ✅ IP addressbased access control lists
- ✅ Track Request IP Address
- ✅ [Use LiteLLM keys/authentication on Pass Through Endpoints](./proxy/pass_through#✨-enterprise---use-litellm-keysauthentication-on-pass-through-endpoints)

View file

@ -17,6 +17,7 @@ Features:
- ✅ [JWT-Auth](../docs/proxy/token_auth.md)
- ✅ [Control available public, private routes (Restrict certain endpoints on proxy)](#control-available-public-private-routes)
- ✅ [Control available public, private routes](#control-available-public-private-routes)
- ✅ [Secret Managers - AWS Key Manager, Google Secret Manager, Azure Key, Hashicorp Vault](../secret)
- ✅ [[BETA] AWS Key Manager v2 - Key Decryption](#beta-aws-key-manager---key-decryption)
- ✅ IP addressbased access control lists
- ✅ Track Request IP Address

View file

@ -3,7 +3,6 @@ import TabItem from '@theme/TabItem';
import Image from '@theme/IdealImage';
# Secret Manager
LiteLLM supports reading secrets from Azure Key Vault, Google Secret Manager
:::info
@ -15,6 +14,8 @@ LiteLLM supports reading secrets from Azure Key Vault, Google Secret Manager
:::
LiteLLM supports **reading secrets (eg. `OPENAI_API_KEY`)** and **writing secrets (eg. Virtual Keys)** from Azure Key Vault, Google Secret Manager, Hashicorp Vault, and AWS Secret Manager.
## Supported Secret Managers
- AWS Key Management Service
@ -23,37 +24,17 @@ LiteLLM supports reading secrets from Azure Key Vault, Google Secret Manager
- [Google Secret Manager](#google-secret-manager)
- Google Key Management Service
- [Hashicorp Vault](#hashicorp-vault)
- [Infisical Secret Manager](#infisical-secret-manager)
- [.env Files](#env-files)
## AWS Key Management V1
:::tip
[BETA] AWS Key Management v2 is on the enterprise tier. Go [here for docs](./proxy/enterprise.md#beta-aws-key-manager---key-decryption)
:::
Use AWS KMS to storing a hashed copy of your Proxy Master Key in the environment.
```bash
export LITELLM_MASTER_KEY="djZ9xjVaZ..." # 👈 ENCRYPTED KEY
export AWS_REGION_NAME="us-west-2"
```
```yaml
general_settings:
key_management_system: "aws_kms"
key_management_settings:
hosted_keys: ["LITELLM_MASTER_KEY"] # 👈 WHICH KEYS ARE STORED ON KMS
```
[**See Decryption Code**](https://github.com/BerriAI/litellm/blob/a2da2a8f168d45648b61279d4795d647d94f90c9/litellm/utils.py#L10182)
## AWS Secret Manager
Store your proxy keys in AWS Secret Manager.
| Feature | Support | Description |
|---------|----------|-------------|
| Reading Secrets | ✅ | Read secrets e.g `OPENAI_API_KEY` |
| Writing Secrets | ✅ | Store secrets e.g `Virtual Keys` |
#### Proxy Usage
1. Save AWS Credentials in your environment
@ -100,6 +81,110 @@ general_settings:
litellm --config /path/to/config.yaml
```
## Hashicorp Vault
| Feature | Support | Description |
|---------|----------|-------------|
| Reading Secrets | ✅ | Read secrets e.g `OPENAI_API_KEY` |
| Writing Secrets | ✅ | Store secrets e.g `Virtual Keys` |
Read secrets from [Hashicorp Vault](https://developer.hashicorp.com/vault/docs/secrets/kv/kv-v2)
**Step 1.** Add Hashicorp Vault details in your environment
LiteLLM supports two methods of authentication:
1. TLS cert authentication - `HCP_VAULT_CLIENT_CERT` and `HCP_VAULT_CLIENT_KEY`
2. Token authentication - `HCP_VAULT_TOKEN`
```bash
HCP_VAULT_ADDR="https://test-cluster-public-vault-0f98180c.e98296b2.z1.hashicorp.cloud:8200"
HCP_VAULT_NAMESPACE="admin"
# Authentication via TLS cert
HCP_VAULT_CLIENT_CERT="path/to/client.pem"
HCP_VAULT_CLIENT_KEY="path/to/client.key"
# OR - Authentication via token
HCP_VAULT_TOKEN="hvs.CAESIG52gL6ljBSdmq*****"
# OPTIONAL
HCP_VAULT_REFRESH_INTERVAL="86400" # defaults to 86400, frequency of cache refresh for Hashicorp Vault
```
**Step 2.** Add to proxy config.yaml
```yaml
general_settings:
key_management_system: "hashicorp_vault"
# [OPTIONAL SETTINGS]
key_management_settings:
store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager
prefix_for_stored_virtual_keys: "litellm/" # OPTIONAL. If set, this prefix will be used for stored virtual keys in the secret manager
access_mode: "read_and_write" # Literal["read_only", "write_only", "read_and_write"]
```
**Step 3.** Start + test proxy
```
$ litellm --config /path/to/config.yaml
```
[Quick Test Proxy](./proxy/user_keys)
#### How it works
**Reading Secrets**
LiteLLM reads secrets from Hashicorp Vault's KV v2 engine using the following URL format:
```
{VAULT_ADDR}/v1/{NAMESPACE}/secret/data/{SECRET_NAME}
```
For example, if you have:
- `HCP_VAULT_ADDR="https://vault.example.com:8200"`
- `HCP_VAULT_NAMESPACE="admin"`
- Secret name: `AZURE_API_KEY`
LiteLLM will look up:
```
https://vault.example.com:8200/v1/admin/secret/data/AZURE_API_KEY
```
#### Expected Secret Format
LiteLLM expects all secrets to be stored as a JSON object with a `key` field containing the secret value.
For example, for `AZURE_API_KEY`, the secret should be stored as:
```json
{
"key": "sk-1234"
}
```
<Image img={require('../img/hcorp.png')} />
**Writing Secrets**
When a Virtual Key is Created / Deleted on LiteLLM, LiteLLM will automatically create / delete the secret in Hashicorp Vault.
- Create Virtual Key on LiteLLM either through the LiteLLM Admin UI or API
<Image img={require('../img/hcorp_create_virtual_key.png')} />
- Check Hashicorp Vault for secret
LiteLLM stores secret under the `prefix_for_stored_virtual_keys` path (default: `litellm/`)
<Image img={require('../img/hcorp_virtual_key.png')} />
## Azure Key Vault
<!--
### Quick Start
@ -240,82 +325,31 @@ $ litellm --test
## .env Files
If no secret manager client is specified, Litellm automatically uses the `.env` file to manage sensitive data. -->
## Hashicorp Vault
## AWS Key Management V1
Read secrets from [Hashicorp Vault](https://developer.hashicorp.com/vault/docs/secrets/kv/kv-v2)
:::tip
Step 1. Add Hashicorp Vault details in your environment
[BETA] AWS Key Management v2 is on the enterprise tier. Go [here for docs](./proxy/enterprise.md#beta-aws-key-manager---key-decryption)
LiteLLM supports two methods of authentication:
:::
1. TLS cert authentication - `HCP_VAULT_CLIENT_CERT` and `HCP_VAULT_CLIENT_KEY`
2. Token authentication - `HCP_VAULT_TOKEN`
Use AWS KMS to storing a hashed copy of your Proxy Master Key in the environment.
```bash
HCP_VAULT_ADDR="https://test-cluster-public-vault-0f98180c.e98296b2.z1.hashicorp.cloud:8200"
HCP_VAULT_NAMESPACE="admin"
# Authentication via TLS cert
HCP_VAULT_CLIENT_CERT="path/to/client.pem"
HCP_VAULT_CLIENT_KEY="path/to/client.key"
# OR - Authentication via token
HCP_VAULT_TOKEN="hvs.CAESIG52gL6ljBSdmq*****"
# OPTIONAL
HCP_VAULT_REFRESH_INTERVAL="86400" # defaults to 86400, frequency of cache refresh for Hashicorp Vault
export LITELLM_MASTER_KEY="djZ9xjVaZ..." # 👈 ENCRYPTED KEY
export AWS_REGION_NAME="us-west-2"
```
Step 2. Add to proxy config.yaml
```yaml
general_settings:
key_management_system: "hashicorp_vault"
key_management_system: "aws_kms"
key_management_settings:
hosted_keys: ["LITELLM_MASTER_KEY"] # 👈 WHICH KEYS ARE STORED ON KMS
```
Step 3. Start + test proxy
[**See Decryption Code**](https://github.com/BerriAI/litellm/blob/a2da2a8f168d45648b61279d4795d647d94f90c9/litellm/utils.py#L10182)
```
$ litellm --config /path/to/config.yaml
```
[Quick Test Proxy](./proxy/user_keys)
#### How it works
LiteLLM reads secrets from Hashicorp Vault's KV v2 engine using the following URL format:
```
{VAULT_ADDR}/v1/{NAMESPACE}/secret/data/{SECRET_NAME}
```
For example, if you have:
- `HCP_VAULT_ADDR="https://vault.example.com:8200"`
- `HCP_VAULT_NAMESPACE="admin"`
- Secret name: `AZURE_API_KEY`
LiteLLM will look up:
```
https://vault.example.com:8200/v1/admin/secret/data/AZURE_API_KEY
```
#### Expected Secret Format
LiteLLM expects all secrets to be stored as a JSON object with a `key` field containing the secret value.
For example, for `AZURE_API_KEY`, the secret should be stored as:
```json
{
"key": "sk-1234"
}
```
<Image img={require('../img/hcorp.png')} />
## All Secret Manager Settings
## **All Secret Manager Settings**
All settings related to secret management

Binary file not shown.

After

Width:  |  Height:  |  Size: 190 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 145 KiB

View file

@ -10,7 +10,6 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
GenerateKeyRequest,
KeyManagementSystem,
KeyRequest,
LiteLLM_AuditLogs,
LiteLLM_VerificationToken,
@ -195,21 +194,28 @@ class KeyManagementEventHooks:
"""
if litellm._key_management_settings is not None:
if litellm._key_management_settings.store_virtual_keys is True:
from litellm.secret_managers.aws_secret_manager_v2 import (
AWSSecretsManagerV2,
from litellm.secret_managers.base_secret_manager import (
BaseSecretManager,
)
# store the key in the secret manager
if (
litellm._key_management_system
== KeyManagementSystem.AWS_SECRET_MANAGER
and isinstance(litellm.secret_manager_client, AWSSecretsManagerV2)
):
if isinstance(litellm.secret_manager_client, BaseSecretManager):
await litellm.secret_manager_client.async_write_secret(
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{secret_name}",
secret_name=KeyManagementEventHooks._get_secret_name(
secret_name
),
secret_value=secret_token,
)
@staticmethod
def _get_secret_name(secret_name: str) -> str:
if litellm._key_management_settings.prefix_for_stored_virtual_keys.endswith(
"/"
):
return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{secret_name}"
else:
return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{secret_name}"
@staticmethod
async def _delete_virtual_keys_from_secret_manager(
keys_being_deleted: List[LiteLLM_VerificationToken],
@ -222,15 +228,17 @@ class KeyManagementEventHooks:
"""
if litellm._key_management_settings is not None:
if litellm._key_management_settings.store_virtual_keys is True:
from litellm.secret_managers.aws_secret_manager_v2 import (
AWSSecretsManagerV2,
from litellm.secret_managers.base_secret_manager import (
BaseSecretManager,
)
if isinstance(litellm.secret_manager_client, AWSSecretsManagerV2):
if isinstance(litellm.secret_manager_client, BaseSecretManager):
for key in keys_being_deleted:
if key.key_alias is not None:
await litellm.secret_manager_client.async_delete_secret(
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key.key_alias}"
secret_name=KeyManagementEventHooks._get_secret_name(
key.key_alias
)
)
else:
verbose_proxy_logger.warning(

View file

@ -5,7 +5,9 @@ model_list:
api_base: https://example-openai-endpoint.onrender.com/chat/completions
api_key: "ishaan"
general_settings:
master_key: sk-1234
alerting: ["slack"]
alerting_threshold: 0.0000001
general_settings:
key_management_system: "hashicorp_vault" # 👈 KEY CHANGE
key_management_settings:
store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager
prefix_for_stored_virtual_keys: "litellm/" # OPTIONAL. If set, this prefix will be used for stored virtual keys in the secret manager
access_mode: "write_only" # Literal["read_only", "write_only", "read_and_write"]

View file

@ -29,8 +29,10 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy._types import KeyManagementSystem
from litellm.types.llms.custom_http import httpxSpecialProvider
from .base_secret_manager import BaseSecretManager
class AWSSecretsManagerV2(BaseAWSLLM):
class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
@classmethod
def validate_environment(cls):
if "AWS_REGION_NAME" not in os.environ:
@ -146,7 +148,6 @@ class AWSSecretsManagerV2(BaseAWSLLM):
secret_name: str,
secret_value: str,
description: Optional[str] = None,
client_request_token: Optional[str] = None,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> dict:
@ -157,7 +158,6 @@ class AWSSecretsManagerV2(BaseAWSLLM):
secret_name: Name of the secret
secret_value: Value to store (can be a JSON string)
description: Optional description for the secret
client_request_token: Optional unique identifier to ensure idempotency
optional_params: Additional AWS parameters
timeout: Request timeout
"""

View file

@ -0,0 +1,95 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Union
import httpx
class BaseSecretManager(ABC):
"""
Abstract base class for secret management implementations.
"""
@abstractmethod
async def async_read_secret(
self,
secret_name: str,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> Optional[str]:
"""
Asynchronously read a secret from the secret manager.
Args:
secret_name (str): Name/path of the secret to read
optional_params (Optional[dict]): Additional parameters specific to the secret manager
timeout (Optional[Union[float, httpx.Timeout]]): Request timeout
Returns:
Optional[str]: The secret value if found, None otherwise
"""
pass
@abstractmethod
def sync_read_secret(
self,
secret_name: str,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> Optional[str]:
"""
Synchronously read a secret from the secret manager.
Args:
secret_name (str): Name/path of the secret to read
optional_params (Optional[dict]): Additional parameters specific to the secret manager
timeout (Optional[Union[float, httpx.Timeout]]): Request timeout
Returns:
Optional[str]: The secret value if found, None otherwise
"""
pass
@abstractmethod
async def async_write_secret(
self,
secret_name: str,
secret_value: str,
description: Optional[str] = None,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> Dict[str, Any]:
"""
Asynchronously write a secret to the secret manager.
Args:
secret_name (str): Name/path of the secret to write
secret_value (str): Value to store
description (Optional[str]): Description of the secret. Some secret managers allow storing a description with the secret.
optional_params (Optional[dict]): Additional parameters specific to the secret manager
timeout (Optional[Union[float, httpx.Timeout]]): Request timeout
Returns:
Dict[str, Any]: Response from the secret manager containing write operation details
"""
pass
@abstractmethod
async def async_delete_secret(
self,
secret_name: str,
recovery_window_in_days: Optional[int] = 7,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> dict:
"""
Async function to delete a secret from the secret manager
Args:
secret_name: Name of the secret to delete
recovery_window_in_days: Number of days before permanent deletion (default: 7)
optional_params: Additional parameters specific to the secret manager
timeout: Request timeout
Returns:
dict: Response from the secret manager containing deletion details
"""
pass

View file

@ -1,5 +1,5 @@
import os
from typing import Optional
from typing import Any, Dict, Optional, Union
import httpx
@ -13,8 +13,10 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._types import KeyManagementSystem
from .base_secret_manager import BaseSecretManager
class HashicorpSecretManager:
class HashicorpSecretManager(BaseSecretManager):
def __init__(self):
from litellm.proxy.proxy_server import CommonProxyErrors, premium_user
@ -110,7 +112,12 @@ class HashicorpSecretManager:
return {"X-Vault-Token": self._auth_via_tls_cert()}
return {"X-Vault-Token": self.vault_token}
async def async_read_secret(self, secret_name: str) -> Optional[str]:
async def async_read_secret(
self,
secret_name: str,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> Optional[str]:
"""
Reads a secret from Vault KV v2 using an async HTTPX client.
secret_name is just the path inside the KV mount (e.g., 'myapp/config').
@ -140,7 +147,12 @@ class HashicorpSecretManager:
verbose_logger.exception(f"Error reading secret from Hashicorp Vault: {e}")
return None
def read_secret(self, secret_name: str) -> Optional[str]:
def sync_read_secret(
self,
secret_name: str,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> Optional[str]:
"""
Reads a secret from Vault KV v2 using a sync HTTPX client.
secret_name is just the path inside the KV mount (e.g., 'myapp/config').
@ -166,6 +178,95 @@ class HashicorpSecretManager:
verbose_logger.exception(f"Error reading secret from Hashicorp Vault: {e}")
return None
async def async_write_secret(
self,
secret_name: str,
secret_value: str,
description: Optional[str] = None,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> Dict[str, Any]:
"""
Writes a secret to Vault KV v2 using an async HTTPX client.
Args:
secret_name: Path inside the KV mount (e.g., 'myapp/config')
secret_value: Value to store
description: Optional description for the secret
optional_params: Additional parameters to include in the secret data
timeout: Request timeout
Returns:
dict: Response containing status and details of the operation
"""
async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"timeout": timeout},
)
try:
url = self.get_url(secret_name)
# Prepare the secret data
data = {"data": {"key": secret_value}}
if description:
data["data"]["description"] = description
response = await async_client.post(
url=url, headers=self._get_request_headers(), json=data
)
response.raise_for_status()
return response.json()
except Exception as e:
verbose_logger.exception(f"Error writing secret to Hashicorp Vault: {e}")
return {"status": "error", "message": str(e)}
async def async_delete_secret(
self,
secret_name: str,
recovery_window_in_days: Optional[int] = 7,
optional_params: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> dict:
"""
Async function to delete a secret from Hashicorp Vault.
In KV v2, this marks the latest version of the secret as deleted.
Args:
secret_name: Name of the secret to delete
recovery_window_in_days: Not used for Vault (Vault handles this internally)
optional_params: Additional parameters specific to the secret manager
timeout: Request timeout
Returns:
dict: Response containing status and details of the operation
"""
async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"timeout": timeout},
)
try:
# For KV v2 delete: /v1/<mount>/data/<path>
url = self.get_url(secret_name)
response = await async_client.delete(
url=url, headers=self._get_request_headers()
)
response.raise_for_status()
# Clear the cache for this secret
self.cache.delete_cache(secret_name)
return {
"status": "success",
"message": f"Secret {secret_name} deleted successfully",
}
except Exception as e:
verbose_logger.exception(f"Error deleting secret from Hashicorp Vault: {e}")
return {"status": "error", "message": str(e)}
def _get_secret_value_from_json_response(
self, json_resp: Optional[dict]
) -> Optional[str]:

View file

@ -291,7 +291,7 @@ def get_secret( # noqa: PLR0915
raise e
elif key_manager == KeyManagementSystem.HASHICORP_VAULT.value:
try:
secret = client.read_secret(secret_name)
secret = client.sync_read_secret(secret_name=secret_name)
if secret is None:
raise ValueError(
f"No secret found in Hashicorp Secret Manager for {secret_name}"

View file

@ -3561,7 +3561,7 @@ async def test_key_generate_with_secret_manager_call(prisma_client):
# read from the secret manager
result = await aws_secret_manager_client.async_read_secret(
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key_alias}"
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{key_alias}"
)
# Assert the correct key is stored in the secret manager
@ -3582,7 +3582,7 @@ async def test_key_generate_with_secret_manager_call(prisma_client):
# Assert the key is deleted from the secret manager
result = await aws_secret_manager_client.async_read_secret(
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key_alias}"
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{key_alias}"
)
assert result is None

View file

@ -13,6 +13,7 @@ sys.path.insert(
from unittest.mock import patch, MagicMock
import logging
from litellm._logging import verbose_logger
import uuid
verbose_logger.setLevel(logging.DEBUG)
@ -42,6 +43,25 @@ mock_vault_response = {
"mount_type": "kv",
}
# Update the mock_vault_response for write operations
mock_write_response = {
"request_id": "80fafb6a-e96a-4c5b-29fa-ff505ac72201",
"lease_id": "",
"renewable": False,
"lease_duration": 0,
"data": {
"created_time": "2025-01-04T16:58:42.684673531Z",
"custom_metadata": None,
"deletion_time": "",
"destroyed": False,
"version": 1,
},
"wrap_info": None,
"warnings": None,
"auth": None,
"mount_type": "kv",
}
def test_hashicorp_secret_manager_get_secret():
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") as mock_get:
@ -52,7 +72,7 @@ def test_hashicorp_secret_manager_get_secret():
mock_get.return_value = mock_response
# Test the secret manager
secret = hashicorp_secret_manager.read_secret("sample-secret-mock")
secret = hashicorp_secret_manager.sync_read_secret("sample-secret-mock")
assert secret == "value-mock"
# Verify the request was made with correct parameters
@ -67,6 +87,84 @@ def test_hashicorp_secret_manager_get_secret():
assert "X-Vault-Token" in mock_get.call_args.kwargs["headers"]
@pytest.mark.asyncio
async def test_hashicorp_secret_manager_write_secret():
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post"
) as mock_post:
# Configure the mock response
mock_response = MagicMock()
mock_response.json.return_value = (
mock_write_response # Use the write-specific response
)
mock_response.raise_for_status.return_value = None
mock_post.return_value = mock_response
# Test the secret manager
secret_name = f"sample-secret-test-{uuid.uuid4()}"
secret_value = f"value-mock-{uuid.uuid4()}"
response = await hashicorp_secret_manager.async_write_secret(
secret_name=secret_name,
secret_value=secret_value,
)
# Verify the response and that the request was made correctly
assert (
response == mock_write_response
) # Compare against write-specific response
mock_post.assert_called_once()
print("CALL ARGS=", mock_post.call_args)
print("call args[1]=", mock_post.call_args[1])
# Verify URL
called_url = mock_post.call_args[1]["url"]
assert secret_name in called_url
assert (
called_url
== f"{hashicorp_secret_manager.vault_addr}/v1/admin/secret/data/{secret_name}"
)
# Verify request body
json_data = mock_post.call_args[1]["json"]
assert "data" in json_data
assert "key" in json_data["data"]
assert json_data["data"]["key"] == secret_value
@pytest.mark.asyncio
async def test_hashicorp_secret_manager_delete_secret():
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete"
) as mock_delete:
# Configure the mock response
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_delete.return_value = mock_response
# Test the secret manager
secret_name = f"sample-secret-test-{uuid.uuid4()}"
response = await hashicorp_secret_manager.async_delete_secret(
secret_name=secret_name
)
# Verify the response
assert response == {
"status": "success",
"message": f"Secret {secret_name} deleted successfully",
}
# Verify the request was made correctly
mock_delete.assert_called_once()
# Verify URL
called_url = mock_delete.call_args[1]["url"]
assert secret_name in called_url
assert (
called_url
== f"{hashicorp_secret_manager.vault_addr}/v1/admin/secret/data/{secret_name}"
)
def test_hashicorp_secret_manager_tls_cert_auth():
with patch("httpx.post") as mock_post:
# Configure the mock response for TLS auth