Spaces:
Configuration error
Configuration error
File size: 7,114 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 |
import os
import sys
import pytest
from dotenv import load_dotenv
load_dotenv()
import os
import httpx
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from unittest.mock import patch, MagicMock
import logging
from litellm._logging import verbose_logger
import uuid
verbose_logger.setLevel(logging.DEBUG)
from litellm.secret_managers.hashicorp_secret_manager import HashicorpSecretManager
hashicorp_secret_manager = HashicorpSecretManager()
mock_vault_response = {
"request_id": "80fafb6a-e96a-4c5b-29fa-ff505ac72201",
"lease_id": "",
"renewable": False,
"lease_duration": 0,
"data": {
"data": {"key": "value-mock"},
"metadata": {
"created_time": "2025-01-01T22:13:50.93942388Z",
"custom_metadata": None,
"deletion_time": "",
"destroyed": False,
"version": 1,
},
},
"wrap_info": None,
"warnings": None,
"auth": None,
"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:
# Configure the mock response using MagicMock
mock_response = MagicMock()
mock_response.json.return_value = mock_vault_response
mock_response.raise_for_status.return_value = None
mock_get.return_value = mock_response
# Test the secret manager
secret = hashicorp_secret_manager.sync_read_secret("sample-secret-mock")
assert secret == "value-mock"
# Verify the request was made with correct parameters
mock_get.assert_called_once()
called_url = mock_get.call_args[0][0]
assert "sample-secret-mock" in called_url
assert (
called_url
== "https://test-cluster-public-vault-0f98180c.e98296b2.z1.hashicorp.cloud:8200/v1/admin/secret/data/sample-secret-mock"
)
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(monkeypatch):
monkeypatch.setenv("HCP_VAULT_TOKEN", "test-client-token-12345")
print("HCP_VAULT_TOKEN=", os.getenv("HCP_VAULT_TOKEN"))
# Mock both httpx.post and httpx.Client
with patch("httpx.Client") as mock_client:
# Configure the mock client and response
mock_response = MagicMock()
mock_response.json.return_value = {
"auth": {
"client_token": "test-client-token-12345",
"lease_duration": 3600,
"renewable": True,
}
}
mock_response.raise_for_status.return_value = None
# Configure the mock client's post method
mock_client_instance = MagicMock()
mock_client_instance.post.return_value = mock_response
mock_client.return_value = mock_client_instance
# Create a new instance with TLS cert config
test_manager = HashicorpSecretManager()
test_manager.tls_cert_path = "cert.pem"
test_manager.tls_key_path = "key.pem"
test_manager.vault_cert_role = "test-role"
test_manager.vault_namespace = "test-namespace"
# Test the TLS auth method
token = test_manager._auth_via_tls_cert()
# Verify the token
assert token == "test-client-token-12345"
# Verify Client was created with correct cert tuple
mock_client.assert_called_once_with(cert=("cert.pem", "key.pem"))
# Verify post was called with correct parameters
mock_client_instance.post.assert_called_once_with(
f"{test_manager.vault_addr}/v1/auth/cert/login",
headers={"X-Vault-Namespace": "test-namespace"},
json={"name": "test-role"},
)
# Verify the token was cached
assert test_manager.cache.get_cache("hcp_vault_token") == "test-client-token-12345"
|