mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-24 23:38:15 +00:00
Krish Dholakia
9b7ebb6a7d
* build(pyproject.toml): add new dev dependencies - for type checking * build: reformat files to fit black * ci: reformat to fit black * ci(test-litellm.yml): make tests run clear * build(pyproject.toml): add ruff * fix: fix ruff checks * build(mypy/): fix mypy linting errors * fix(hashicorp_secret_manager.py): fix passing cert for tls auth * build(mypy/): resolve all mypy errors * test: update test * fix: fix black formatting * build(pre-commit-config.yaml): use poetry run black * fix(proxy_server.py): fix linting error * fix: fix ruff safe representation error
214 lines
6.9 KiB
Python
214 lines
6.9 KiB
Python
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"
|