diff --git a/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py new file mode 100644 index 00000000000..f860fecf47d --- /dev/null +++ b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py @@ -0,0 +1,65 @@ +import hashlib +import json +import uuid +from typing import Final, Protocol, cast + +import pytest +from e2e_metadata import Domain, Subject, meta +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 +from litellm.secret_managers.main import get_secret + +pytestmark = pytest.mark.e2e + +_AWS_FIXTURE_MASTER_KEY_SHA256: Final = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STRING_MAP: Final = TypeAdapter(dict[str, str]) +_STRING: Final = TypeAdapter(str) + + +class _AWSSecretManagerOperations(Protocol): + async def async_write_secret(self, secret_name: str, secret_value: str) -> object: ... + + async def async_delete_secret( + self, secret_name: str, recovery_window_in_days: int | None = 7 + ) -> object: ... + + +@pytest.mark.asyncio +@meta(Subject(domain=Domain.DEPLOY_OPS)) +async def test_aws_secret_manager() -> None: + AWSSecretsManagerV2.load_aws_secret_manager(use_aws_secret_manager=True) + + fixture_value: Final[str | None] = cast( + str | None, TypeAdapter(str | None).validate_python(get_secret("litellm_master_key")) + ) + if fixture_value is not None: + fixture_data: Final = _STRING_MAP.validate_json(_STRING.validate_python(fixture_value)) + assert hashlib.sha256(fixture_data["litellm_master_key"].encode()).hexdigest() == ( + _AWS_FIXTURE_MASTER_KEY_SHA256 + ) + return + + secret_name: Final = f"litellm-e2e-secret-manager-{uuid.uuid4().hex}" + secret_value: Final = f"sk-litellm-e2e-{uuid.uuid4().hex}" + manager: Final = litellm.secret_manager_client + assert isinstance(manager, AWSSecretsManagerV2) + secret_manager: Final = cast(_AWSSecretManagerOperations, manager) + try: + creation: Final = _JSON_OBJECT.validate_python( + await secret_manager.async_write_secret( + secret_name, + json.dumps({"litellm_master_key": secret_value}), + ) + ) + assert creation.get("Name") == secret_name + created_value: Final = _STRING.validate_python(get_secret(secret_name)) + created_data: Final = _STRING_MAP.validate_json(created_value) + assert created_data["litellm_master_key"] == secret_value + finally: + deletion: Final = _JSON_OBJECT.validate_python( + await secret_manager.async_delete_secret(secret_name, recovery_window_in_days=7) + ) + assert deletion.get("Name") == secret_name diff --git a/tests/integration/providers/test_google_genai_proxy_sdk_wire.py b/tests/integration/providers/test_google_genai_proxy_sdk_wire.py new file mode 100644 index 00000000000..2e5856a5923 --- /dev/null +++ b/tests/integration/providers/test_google_genai_proxy_sdk_wire.py @@ -0,0 +1,175 @@ +import json +from collections.abc import Callable, Mapping +from typing import Final, Literal, TypeAlias + +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from google import genai +from google.genai import types +from integration._support.client import Gateway, Scenario +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "gemini-2.5-flash" +_API_KEY: Final = "synthetic-gemini-key" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_MODEL_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models/{_BACKEND}" +_PROMPT: Final = "Reply with only the single word: pong" +_PROVIDERS: Final[tuple[Literal["gemini", "vertex_ai"], ...]] = ("gemini", "vertex_ai") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +Provider: TypeAlias = Literal["gemini", "vertex_ai"] + + +def _service_account_json(token_url: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": _PROJECT, + "private_key_id": "scripted", + "private_key": private_key, + "client_email": f"scripted@{_PROJECT}.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{token_url}/_oauth/authorize", + "token_uri": f"{token_url}/_oauth/token", + } + ) + + +def _register_model(gateway: Gateway, scenario: Scenario, provider: Provider, wire: Wire) -> str: + if provider == "gemini": + return scenario.model(model=f"gemini/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + return scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=f"{wire.url}{_MODEL_PATH}", + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=_service_account_json(wire.url), + ) + + +def _target(provider: Provider, stream: bool) -> str: + path: Final = f"/models/{_BACKEND}" if provider == "gemini" else _MODEL_PATH + if stream: + return f"{path}:streamGenerateContent?alt=sse" + return f"{path}:generateContent" + + +def _response_body() -> Mapping[str, JsonValue]: + return { + "responseId": "scripted-response", + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "pong"}]}, + "index": 0, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 1, "totalTokenCount": 10}, + "modelVersion": _BACKEND, + } + + +def _expected_request_body(provider: Provider) -> Mapping[str, JsonValue]: + request: Final = { + "model": _BACKEND, + "contents": [{"parts": [{"text": _PROMPT}], "role": "user"}], + } + if provider == "gemini": + return { + **request, + "tools": None, + "generationConfig": {"temperature": 0, "topP": 0.95, "topK": 20}, + } + return { + **request, + "generationConfig": {"temperature": 0, "top_p": 0.95, "top_k": 20}, + } + + +def _respond(provider: Provider, stream: bool) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + body=b'{"access_token":"scripted-token","expires_in":3600,"token_type":"Bearer"}' + ) + assert request.target == _target(provider, stream) + if provider == "gemini": + assert request.headers["x-goog-api-key"] == _API_KEY + else: + assert request.headers["authorization"] == "Bearer scripted-token" + assert _JSON_OBJECT.validate_json(request.body) == _expected_request_body(provider) + if stream: + return Reply( + content_type="text/event-stream", + chunks=(f"data: {json.dumps(_response_body())}\n\n".encode(),), + ) + return Reply(body=json.dumps(_response_body()).encode()) + + return respond + + +def _client(gateway: Gateway) -> genai.Client: + return genai.Client( + api_key=gateway.key, + http_options={"base_url": str(gateway.client.base_url)}, + ) + + +@pytest.mark.parametrize("provider", _PROVIDERS) +def test_google_genai_sdk_non_streaming_forwards_body_and_response( + gateway: Gateway, provider: Provider +) -> None: + with wire_server(_respond(provider, stream=False)) as wire, gateway.scenario() as scenario: + model: Final = _register_model(gateway, scenario, provider, wire) + response: Final = _client(gateway).models.generate_content( + model=model, + contents=types.Part.from_text(text=_PROMPT), + config=types.GenerateContentConfig(temperature=0, top_p=0.95, top_k=20), + ) + assert response.text == "pong" + assert response.usage_metadata.prompt_token_count == 9 + assert response.usage_metadata.candidates_token_count == 1 + + +@pytest.mark.parametrize("provider", _PROVIDERS) +def test_google_genai_sdk_streaming_forwards_body_and_response( + gateway: Gateway, provider: Provider +) -> None: + with wire_server(_respond(provider, stream=True)) as wire, gateway.scenario() as scenario: + model: Final = _register_model(gateway, scenario, provider, wire) + chunks: Final = tuple( + _client(gateway).models.generate_content_stream( + model=model, + contents=types.Part.from_text(text=_PROMPT), + config=types.GenerateContentConfig(temperature=0, top_p=0.95, top_k=20), + ) + ) + assert tuple(chunk.text for chunk in chunks) == ("pong",) + + +@pytest.mark.parametrize("provider", _PROVIDERS) +def test_google_genai_sdk_dict_streaming_forwards_body_and_response( + gateway: Gateway, provider: Provider +) -> None: + with wire_server(_respond(provider, stream=True)) as wire, gateway.scenario() as scenario: + model: Final = _register_model(gateway, scenario, provider, wire) + chunks: Final = tuple( + _client(gateway).models.generate_content_stream( + model=model, + contents={"text": _PROMPT}, + config={"temperature": 0, "top_p": 0.95, "top_k": 20}, + ) + ) + assert tuple(chunk.text for chunk in chunks) == ("pong",) diff --git a/tests/litellm_utils_tests/base_token_counter_test.py b/tests/litellm_utils_tests/base_token_counter_test.py deleted file mode 100644 index f4c90118a03..00000000000 --- a/tests/litellm_utils_tests/base_token_counter_test.py +++ /dev/null @@ -1,72 +0,0 @@ -""" -Base Token Counter Test Suite. - -This module provides an abstract base test class that enforces common tests -across all token counter implementations. Similar to base_llm_unit_tests.py -for LLM chat tests. - -Usage: - Create a test class that inherits from BaseTokenCounterTest and implement - the abstract methods to provide provider-specific configuration. -""" - -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional - -import pytest - - -from litellm.llms.base_llm.base_utils import BaseTokenCounter - - -class BaseTokenCounterTest(ABC): - """ - Abstract base test class for token counter implementations. - - Subclasses must implement: - - get_token_counter(): Returns the token counter instance - - get_test_model(): Returns the model name to use for testing - - get_test_messages(): Returns test messages for token counting - - get_deployment_config(): Returns deployment configuration with credentials - - get_custom_llm_provider(): Returns the provider name for should_use_token_counting_api - """ - - @abstractmethod - def get_token_counter(self) -> BaseTokenCounter: - """Must return the token counter instance to test.""" - pass - - @abstractmethod - def get_test_model(self) -> str: - """Must return the model name to use for testing.""" - pass - - @abstractmethod - def get_test_messages(self) -> List[Dict[str, Any]]: - """Must return test messages for token counting.""" - pass - - @abstractmethod - def get_deployment_config(self) -> Dict[str, Any]: - """Must return deployment configuration with credentials.""" - pass - - @abstractmethod - def get_custom_llm_provider(self) -> str: - """Must return the provider name for should_use_token_counting_api check.""" - pass - - @pytest.fixture(autouse=True) - def _handle_missing_credentials(self): - """Fixture to skip tests when credentials are missing.""" - try: - yield - except Exception as e: - error_str = str(e).lower() - if ( - "api key" in error_str - or "api_key" in error_str - or "unauthorized" in error_str - ): - pytest.skip(f"Missing or invalid credentials: {e}") - raise diff --git a/tests/litellm_utils_tests/conftest.py b/tests/litellm_utils_tests/conftest.py deleted file mode 100644 index 002ed594d3f..00000000000 --- a/tests/litellm_utils_tests/conftest.py +++ /dev/null @@ -1,106 +0,0 @@ -# conftest.py - -import asyncio -import importlib - -import pytest - -import litellm # noqa: E402,F401 - -from tests._vcr_conftest_common import ( # noqa: E402,F401 - VerboseReporterState, - _pin_multipart_boundary, - apply_vcr_auto_marker_to_items, - emit_cassette_cache_session_banner, - emit_vcr_classification_summary, - emit_vcr_diagnostic_log, - install_live_call_probe, - record_vcr_outcome, - register_persister_if_enabled, - reset_vcr_diag_dir, - vcr_config_dict, -) - -_verbose_state = VerboseReporterState() - -_VCR_INCOMPATIBLE_FILES = frozenset() - -_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = () - - -@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) - print(litellm) - yield - - # Teardown code (executes after the yield point) - loop.close() # Close the loop created earlier - asyncio.set_event_loop(None) # Remove the reference to the loop - - -@pytest.fixture(scope="module") -def vcr_config(): - return vcr_config_dict() - - -def pytest_recording_configure(config, vcr): - register_persister_if_enabled(vcr) - - -@pytest.hookimpl(hookwrapper=True) -def pytest_runtest_makereport(item, call): - outcome = yield - rep = outcome.get_result() - setattr(item, f"rep_{rep.when}", rep) - - -@pytest.fixture(autouse=True) -def _vcr_outcome_gate(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) - - -def pytest_configure(config): - _verbose_state.remember_pluginmanager(config) - reset_vcr_diag_dir() - - -def pytest_runtest_logreport(report): - _verbose_state.maybe_emit_verdict(report) - - -def pytest_collection_modifyitems(config, items): - apply_vcr_auto_marker_to_items( - items, - skip_files=_VCR_INCOMPATIBLE_FILES, - skip_nodeid_suffixes=_VCR_INCOMPATIBLE_NODEID_SUFFIXES, - ) - - # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests - custom_logger_tests = [ - item for item in items if "custom_logger" in item.parent.name - ] - other_tests = [item for item in items if "custom_logger" not in item.parent.name] - - # Sort tests based on their names - custom_logger_tests.sort(key=lambda x: x.name) - other_tests.sort(key=lambda x: x.name) - - # Reorder the items list - items[:] = custom_logger_tests + other_tests - - -def pytest_terminal_summary(terminalreporter, exitstatus, config): - emit_cassette_cache_session_banner(terminalreporter) - emit_vcr_classification_summary(terminalreporter) - emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/litellm_utils_tests/log.txt b/tests/litellm_utils_tests/log.txt deleted file mode 100644 index bce7b79031e..00000000000 --- a/tests/litellm_utils_tests/log.txt +++ /dev/null @@ -1,24 +0,0 @@ -============================= test session starts ============================== -platform darwin -- Python 3.13.1, pytest-8.3.5, pluggy-1.5.0 -- /Users/krrishdholakia/Documents/litellm/myenv/bin/python3.13 -cachedir: .pytest_cache -rootdir: /Users/krrishdholakia/Documents/litellm -configfile: pyproject.toml -plugins: respx-0.22.0, postgresql-7.0.1, anyio-4.4.0, asyncio-0.26.0, mock-3.14.0, ddtrace-2.19.0rc1 -asyncio: mode=Mode.STRICT, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function -collecting ... collected 3 items - -test_supports_tool_choice.py::test_check_provider_match PASSED [ 33%] -test_supports_tool_choice.py::test_supports_tool_choice PASSED [ 66%] -test_supports_tool_choice.py::test_supports_tool_choice_simple_tests PASSED [100%] - -=============================== warnings summary =============================== -../../myenv/lib/python3.13/site-packages/pydantic/_internal/_config.py:295 - /Users/krrishdholakia/Documents/litellm/myenv/lib/python3.13/site-packages/pydantic/_internal/_config.py:295: PydanticDeprecatedSince20: Support for class-based `config` is deprecated, use ConfigDict instead. Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.10/migration/ - warnings.warn(DEPRECATION_MESSAGE, DeprecationWarning) - -../../litellm/caching/llm_caching_handler.py:17 - /Users/krrishdholakia/Documents/litellm/litellm/caching/llm_caching_handler.py:17: DeprecationWarning: There is no current event loop - event_loop = asyncio.get_event_loop() - --- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html -======================== 3 passed, 2 warnings in 0.92s ========================= diff --git a/tests/litellm_utils_tests/test_anthropic_token_counter.py b/tests/litellm_utils_tests/test_anthropic_token_counter.py deleted file mode 100644 index df3d198b6cf..00000000000 --- a/tests/litellm_utils_tests/test_anthropic_token_counter.py +++ /dev/null @@ -1,41 +0,0 @@ -""" -Anthropic Token Counter Tests. - -Tests for the Anthropic token counter implementation using the base test suite. -""" - -import os -from typing import Any, Dict, List - -import pytest - - -from litellm.llms.anthropic.count_tokens import AnthropicTokenCounter -from litellm.llms.base_llm.base_utils import BaseTokenCounter -from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest - - -class TestAnthropicTokenCounter(BaseTokenCounterTest): - """Test suite for Anthropic token counter.""" - - def get_token_counter(self) -> BaseTokenCounter: - return AnthropicTokenCounter() - - def get_test_model(self) -> str: - return "claude-haiku-4-5-20251001" - - def get_test_messages(self) -> List[Dict[str, Any]]: - return [{"role": "user", "content": "Hello, how are you today?"}] - - def get_deployment_config(self) -> Dict[str, Any]: - api_key = os.getenv("ANTHROPIC_API_KEY") - if not api_key: - pytest.skip("ANTHROPIC_API_KEY not set") - return { - "litellm_params": { - "api_key": api_key, - } - } - - def get_custom_llm_provider(self) -> str: - return "anthropic" diff --git a/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py b/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py deleted file mode 100644 index 50631eb9341..00000000000 --- a/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Azure AI Anthropic Token Counter Tests. - -Tests for the Azure AI Anthropic token counter implementation using the base test suite. -""" - -import os -from typing import Any, Dict, List - -import pytest - - -from litellm.llms.azure_ai.anthropic.count_tokens import AzureAIAnthropicTokenCounter -from litellm.llms.base_llm.base_utils import BaseTokenCounter -from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest - - -class TestAzureAIAnthropicTokenCounter(BaseTokenCounterTest): - """Test suite for Azure AI Anthropic token counter.""" - - def get_token_counter(self) -> BaseTokenCounter: - return AzureAIAnthropicTokenCounter() - - def get_test_model(self) -> str: - return "claude-sonnet-4-6" - - def get_test_messages(self) -> List[Dict[str, Any]]: - return [{"role": "user", "content": "Hello, how are you today?"}] - - def get_deployment_config(self) -> Dict[str, Any]: - api_key = os.getenv("AZURE_ANTHROPIC_API_KEY") - api_base = os.getenv("AZURE_AI_SWEDEN_API_BASE") - - if not api_key: - pytest.skip("AZURE_AI_API_KEY not set") - if not api_base: - pytest.skip("AZURE_AI_API_BASE not set") - - return { - "litellm_params": { - "api_key": api_key, - "api_base": api_base, - } - } - - def get_custom_llm_provider(self) -> str: - return "azure_ai" diff --git a/tests/litellm_utils_tests/test_hashicorp.py b/tests/litellm_utils_tests/test_hashicorp.py deleted file mode 100644 index ac9d4af3f53..00000000000 --- a/tests/litellm_utils_tests/test_hashicorp.py +++ /dev/null @@ -1,844 +0,0 @@ -import os -import pytest -from dotenv import load_dotenv - -load_dotenv() -import httpx - -from unittest.mock import patch, MagicMock -import logging -from litellm._logging import verbose_logger -from litellm._uuid import uuid - -verbose_logger.setLevel(logging.DEBUG) - -# Minimal setup for module-level instantiation -import litellm.proxy.proxy_server - -litellm.proxy.proxy_server.premium_user = True - -from litellm.secret_managers.hashicorp_secret_manager import HashicorpSecretManager - - -@pytest.fixture -def hashicorp_secret_manager(): - """Provide a fresh HashicorpSecretManager per test to avoid shared state.""" - has_token = bool(os.getenv("HCP_VAULT_TOKEN")) - has_approle = bool( - os.getenv("HCP_VAULT_APPROLE_ROLE_ID") - and os.getenv("HCP_VAULT_APPROLE_SECRET_ID") - ) - has_tls_cert = bool( - os.getenv("HCP_VAULT_CLIENT_CERT") and os.getenv("HCP_VAULT_CLIENT_KEY") - ) - if not (has_token or has_approle or has_tls_cert): - pytest.skip( - "Skipping Hashicorp tests: set HCP_VAULT_TOKEN, AppRole vars, or TLS cert vars." - ) - - manager = HashicorpSecretManager() - manager.vault_addr = ( - "https://test-cluster-public-vault-0f98180c.e98296b2.z1.hashicorp.cloud:8200" - ) - manager.vault_namespace = "admin" - manager.vault_mount_name = "secret" - manager.vault_path_prefix = None - return manager - - -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(hashicorp_secret_manager): - 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(hashicorp_secret_manager): - 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_write_secret_with_team_overrides( - hashicorp_secret_manager, -): - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" - ) as mock_post: - mock_response = MagicMock() - mock_response.json.return_value = mock_write_response - mock_response.raise_for_status.return_value = None - mock_post.return_value = mock_response - - secret_value = "value-mock" - team_settings = { - "namespace": "team-namespace", - "mount": "kv-team", - "path_prefix": "teams/custom", - "data": "password", - } - - response = await hashicorp_secret_manager.async_write_secret( - secret_name="team-secret", - secret_value=secret_value, - optional_params=team_settings, - ) - - assert response == mock_write_response - mock_post.assert_called_once() - - called_url = mock_post.call_args[1]["url"] - expected_url = ( - f"{hashicorp_secret_manager.vault_addr}/v1/" - "team-namespace/kv-team/data/teams/custom/team-secret" - ) - assert called_url == expected_url - - json_data = mock_post.call_args[1]["json"] - assert json_data["data"] == {"password": secret_value} - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_delete_secret(hashicorp_secret_manager): - 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}" - ) - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_delete_secret_with_team_overrides( - hashicorp_secret_manager, -): - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete" - ) as mock_delete: - mock_response = MagicMock() - mock_response.raise_for_status.return_value = None - mock_delete.return_value = mock_response - - team_settings = { - "namespace": "team-namespace", - "mount": "kv-team", - "path_prefix": "teams/custom", - } - - response = await hashicorp_secret_manager.async_delete_secret( - secret_name="team-secret", optional_params=team_settings - ) - - assert response == { - "status": "success", - "message": "Secret team-secret deleted successfully", - } - - mock_delete.assert_called_once() - called_url = mock_delete.call_args[1]["url"] - expected_url = ( - f"{hashicorp_secret_manager.vault_addr}/v1/" - "team-namespace/kv-team/data/teams/custom/team-secret" - ) - assert called_url == expected_url - - -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" - ) - - -def test_hashicorp_secret_manager_approle_auth(monkeypatch): - """ - Test AppRole authentication makes the expected POST request to the correct URL. - """ - monkeypatch.setenv("HCP_VAULT_TOKEN", "test-token-12345") - - with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: - mock_response = MagicMock() - mock_response.json.return_value = { - "auth": { - "client_token": "hvs.approle-token-67890", - "lease_duration": 2764800, - "renewable": True, - } - } - mock_response.raise_for_status.return_value = None - mock_post.return_value = mock_response - - test_manager = HashicorpSecretManager() - test_manager.approle_role_id = "test-role-id-123" - test_manager.approle_secret_id = "test-secret-id-456" - test_manager.approle_mount_path = "approle" - - token = test_manager._auth_via_approle() - - assert token == "hvs.approle-token-67890" - - expected_headers = {} - if test_manager.vault_namespace: - expected_headers["X-Vault-Namespace"] = test_manager.vault_namespace - - mock_post.assert_called_once_with( - url=f"{test_manager.vault_addr}/v1/auth/approle/login", - headers=expected_headers, - json={ - "role_id": "test-role-id-123", - "secret_id": "test-secret-id-456", - }, - ) - - assert ( - test_manager.cache.get_cache("hcp_vault_approle_token") - == "hvs.approle-token-67890" - ) - - -def test_hashicorp_custom_mount_and_prefix(hashicorp_secret_manager): - """Test URL construction with custom mount name and path prefix using get_url method.""" - # Save original values - original_mount = hashicorp_secret_manager.vault_mount_name - original_prefix = hashicorp_secret_manager.vault_path_prefix - original_namespace = hashicorp_secret_manager.vault_namespace - - try: - # Test that existing manager uses default "secret" mount and namespace "admin" - url = hashicorp_secret_manager.get_url("my-secret") - assert "/secret/data/" in url - assert "my-secret" in url - - # Test custom mount name - hashicorp_secret_manager.vault_mount_name = "kv" - hashicorp_secret_manager.vault_path_prefix = None - hashicorp_secret_manager.vault_namespace = None - url = hashicorp_secret_manager.get_url("my-secret") - assert url == f"{hashicorp_secret_manager.vault_addr}/v1/kv/data/my-secret" - - # Test path prefix - hashicorp_secret_manager.vault_mount_name = "secret" - hashicorp_secret_manager.vault_path_prefix = "myapp" - url = hashicorp_secret_manager.get_url("my-secret") - assert ( - url - == f"{hashicorp_secret_manager.vault_addr}/v1/secret/data/myapp/my-secret" - ) - - # Test both custom mount and prefix - hashicorp_secret_manager.vault_mount_name = "kv" - hashicorp_secret_manager.vault_path_prefix = "production" - url = hashicorp_secret_manager.get_url("my-secret") - assert ( - url - == f"{hashicorp_secret_manager.vault_addr}/v1/kv/data/production/my-secret" - ) - finally: - # Restore original values - hashicorp_secret_manager.vault_mount_name = original_mount - hashicorp_secret_manager.vault_path_prefix = original_prefix - hashicorp_secret_manager.vault_namespace = original_namespace - - -@pytest.mark.parametrize( - "malicious_secret_name", - [ - "../../../other-app/creds", - "litellm/../../secret", - "foo\nbar", - "foo
bar", - "foo
bar", - "foo\x85bar", - ], -) -def test_hashicorp_get_url_rejects_path_traversal(monkeypatch, malicious_secret_name): - """ - Regression test: get_url must reject an invalid secret_name instead of - building a URL from it. - - Uses monkeypatch + a directly-constructed manager (not the shared - hashicorp_secret_manager fixture) so this runs in CI without real Vault - credentials configured; get_url performs no I/O. - """ - monkeypatch.setenv("HCP_VAULT_TOKEN", "test-token-for-get-url-only") - manager = HashicorpSecretManager() - - with pytest.raises(ValueError, match='Invalid secret_name'): - manager.get_url(malicious_secret_name) - - -mock_old_vault_response = { - "request_id": "80fafb6a-e96a-4c5b-29fa-ff505ac72201", - "lease_id": "", - "renewable": False, - "lease_duration": 0, - "data": { - "data": {"key": "old-secret-value"}, - "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", -} - -mock_new_vault_response = { - "request_id": "80fafb6a-e96a-4c5b-29fa-ff505ac72201", - "lease_id": "", - "renewable": False, - "lease_duration": 0, - "data": { - "data": {"key": "new-secret-value"}, - "metadata": { - "created_time": "2025-01-02T22:13:50.93942388Z", - "custom_metadata": None, - "deletion_time": "", - "destroyed": False, - "version": 1, - }, - }, - "wrap_info": None, - "warnings": None, - "auth": None, - "mount_type": "kv", -} - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_rotate_secret_different_names( - hashicorp_secret_manager, -): - """Test rotating a secret with different names (create new, delete old).""" - with ( - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" - ) as mock_post, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete" - ) as mock_delete, - ): - # Mock GET for current secret check - mock_get_response_current = MagicMock() - mock_get_response_current.json.return_value = mock_old_vault_response - mock_get_response_current.raise_for_status.return_value = None - - # Mock POST for creating new secret - mock_post_response = MagicMock() - mock_post_response.json.return_value = mock_write_response - mock_post_response.raise_for_status.return_value = None - - # Mock GET for verifying new secret - mock_get_response_new = MagicMock() - mock_get_response_new.json.return_value = mock_new_vault_response - mock_get_response_new.raise_for_status.return_value = None - - # Mock DELETE for deleting old secret - mock_delete_response = MagicMock() - mock_delete_response.raise_for_status.return_value = None - - # Configure mock return values - mock_get.side_effect = [mock_get_response_current, mock_get_response_new] - mock_post.return_value = mock_post_response - mock_delete.return_value = mock_delete_response - - current_secret_name = f"old-secret-{uuid.uuid4()}" - new_secret_name = f"new-secret-{uuid.uuid4()}" - new_secret_value = "new-secret-value" - - response = await hashicorp_secret_manager.async_rotate_secret( - current_secret_name=current_secret_name, - new_secret_name=new_secret_name, - new_secret_value=new_secret_value, - ) - - # Verify response - assert response == mock_write_response - - # Verify GET was called twice (check current, verify new) - assert mock_get.call_count == 2 - - # Verify POST was called once (create new secret) - mock_post.assert_called_once() - - # Verify DELETE was called once (delete old secret) - mock_delete.assert_called_once() - - # Verify URLs - get_calls = mock_get.call_args_list - assert current_secret_name in get_calls[0][1]["url"] - assert new_secret_name in get_calls[1][1]["url"] - - delete_url = mock_delete.call_args[1]["url"] - assert current_secret_name in delete_url - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_rotate_secret_same_name( - hashicorp_secret_manager, -): - """Test rotating a secret with the same name (update value only, no delete).""" - with ( - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" - ) as mock_post, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete" - ) as mock_delete, - ): - # Mock GET for current secret check - mock_get_response_current = MagicMock() - mock_get_response_current.json.return_value = mock_old_vault_response - mock_get_response_current.raise_for_status.return_value = None - - # Mock POST for updating secret - mock_post_response = MagicMock() - mock_post_response.json.return_value = mock_write_response - mock_post_response.raise_for_status.return_value = None - - # Mock GET for verifying updated secret - use updated value - mock_get_response_new = MagicMock() - mock_updated_vault_response = { - "request_id": "80fafb6a-e96a-4c5b-29fa-ff505ac72201", - "lease_id": "", - "renewable": False, - "lease_duration": 0, - "data": { - "data": {"key": "updated-secret-value"}, - "metadata": { - "created_time": "2025-01-02T22:13:50.93942388Z", - "custom_metadata": None, - "deletion_time": "", - "destroyed": False, - "version": 1, - }, - }, - "wrap_info": None, - "warnings": None, - "auth": None, - "mount_type": "kv", - } - mock_get_response_new.json.return_value = mock_updated_vault_response - mock_get_response_new.raise_for_status.return_value = None - - # Configure mock return values - mock_get.side_effect = [mock_get_response_current, mock_get_response_new] - mock_post.return_value = mock_post_response - - secret_name = f"same-secret-{uuid.uuid4()}" - new_secret_value = "updated-secret-value" - - response = await hashicorp_secret_manager.async_rotate_secret( - current_secret_name=secret_name, - new_secret_name=secret_name, # Same name - new_secret_value=new_secret_value, - ) - - # Verify response - assert response == mock_write_response - - # Verify GET was called twice (check current, verify new) - assert mock_get.call_count == 2 - - # Verify POST was called once (update secret) - mock_post.assert_called_once() - - # Verify DELETE was NOT called (same name means no delete) - mock_delete.assert_not_called() - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_rotate_secret_current_not_found( - hashicorp_secret_manager, -): - """Test rotating a secret when current secret doesn't exist.""" - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get: - # Mock GET to return 404 - mock_404_response = MagicMock() - mock_404_response.status_code = 404 - mock_404_response.text = "Not Found" - - http_error = httpx.HTTPStatusError( - "Not Found", - request=MagicMock(), - response=mock_404_response, - ) - mock_get.side_effect = http_error - - current_secret_name = f"non-existent-{uuid.uuid4()}" - new_secret_name = f"new-secret-{uuid.uuid4()}" - - response = await hashicorp_secret_manager.async_rotate_secret( - current_secret_name=current_secret_name, - new_secret_name=new_secret_name, - new_secret_value="new-value", - ) - - # Verify error response - assert response["status"] == "error" - assert current_secret_name in response["message"] - assert "not found" in response["message"].lower() - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_rotate_secret_write_fails( - hashicorp_secret_manager, -): - """Test rotating a secret when write fails.""" - with ( - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" - ) as mock_post, - ): - # Mock GET for current secret check - mock_get_response_current = MagicMock() - mock_get_response_current.json.return_value = mock_old_vault_response - mock_get_response_current.raise_for_status.return_value = None - mock_get.return_value = mock_get_response_current - - # Mock POST to return error - mock_post_response = MagicMock() - mock_post_response.json.return_value = { - "status": "error", - "message": "Write failed", - } - mock_post.return_value = mock_post_response - - current_secret_name = f"old-secret-{uuid.uuid4()}" - new_secret_name = f"new-secret-{uuid.uuid4()}" - - response = await hashicorp_secret_manager.async_rotate_secret( - current_secret_name=current_secret_name, - new_secret_name=new_secret_name, - new_secret_value="new-value", - ) - - # Verify error response - assert response["status"] == "error" - assert "Write failed" in response["message"] - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_rotate_secret_with_team_overrides( - hashicorp_secret_manager, -): - """Test rotating a secret with optional_params (team settings).""" - with ( - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" - ) as mock_post, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete" - ) as mock_delete, - ): - # Mock GET for current secret check - mock_get_response_current = MagicMock() - mock_get_response_current.json.return_value = mock_old_vault_response - mock_get_response_current.raise_for_status.return_value = None - - # Mock POST for creating new secret - mock_post_response = MagicMock() - mock_post_response.json.return_value = mock_write_response - mock_post_response.raise_for_status.return_value = None - - # Mock GET for verifying new secret - use password key for team settings - mock_get_response_new = MagicMock() - mock_team_vault_response = { - "request_id": "80fafb6a-e96a-4c5b-29fa-ff505ac72201", - "lease_id": "", - "renewable": False, - "lease_duration": 0, - "data": { - "data": {"password": "new-team-secret-value"}, # Use password key - "metadata": { - "created_time": "2025-01-02T22:13:50.93942388Z", - "custom_metadata": None, - "deletion_time": "", - "destroyed": False, - "version": 1, - }, - }, - "wrap_info": None, - "warnings": None, - "auth": None, - "mount_type": "kv", - } - mock_get_response_new.json.return_value = mock_team_vault_response - mock_get_response_new.raise_for_status.return_value = None - - # Mock DELETE for deleting old secret - mock_delete_response = MagicMock() - mock_delete_response.raise_for_status.return_value = None - - # Configure mock return values - mock_get.side_effect = [mock_get_response_current, mock_get_response_new] - mock_post.return_value = mock_post_response - mock_delete.return_value = mock_delete_response - - team_settings = { - "secret_manager_settings": { - "namespace": "team-namespace", - "mount": "kv-team", - "path_prefix": "teams/custom", - "data": "password", - } - } - - current_secret_name = "team-old-secret" - new_secret_name = "team-new-secret" - new_secret_value = "new-team-secret-value" - - response = await hashicorp_secret_manager.async_rotate_secret( - current_secret_name=current_secret_name, - new_secret_name=new_secret_name, - new_secret_value=new_secret_value, - optional_params=team_settings, - ) - - # Verify response - assert response == mock_write_response - - # Verify URLs use team settings - get_calls = mock_get.call_args_list - assert "team-namespace" in get_calls[0][1]["url"] - assert "kv-team" in get_calls[0][1]["url"] - assert "teams/custom" in get_calls[0][1]["url"] - - delete_url = mock_delete.call_args[1]["url"] - assert "team-namespace" in delete_url - assert "kv-team" in delete_url - - -@pytest.mark.asyncio -async def test_hashicorp_secret_manager_rotate_secret_value_mismatch( - hashicorp_secret_manager, -): - """Test rotating a secret when verification shows value mismatch.""" - with ( - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get, - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" - ) as mock_post, - ): - # Mock GET for current secret check - mock_get_response_current = MagicMock() - mock_get_response_current.json.return_value = mock_old_vault_response - mock_get_response_current.raise_for_status.return_value = None - - # Mock POST for creating new secret - mock_post_response = MagicMock() - mock_post_response.json.return_value = mock_write_response - mock_post_response.raise_for_status.return_value = None - - # Mock GET for verifying new secret - return different value - mock_get_response_new = MagicMock() - mock_get_response_new.json.return_value = { - "data": { - "data": {"key": "different-value"}, # Different from expected - } - } - mock_get_response_new.raise_for_status.return_value = None - - # Configure mock return values - mock_get.side_effect = [mock_get_response_current, mock_get_response_new] - mock_post.return_value = mock_post_response - - current_secret_name = f"old-secret-{uuid.uuid4()}" - new_secret_name = f"new-secret-{uuid.uuid4()}" - new_secret_value = "expected-value" - - response = await hashicorp_secret_manager.async_rotate_secret( - current_secret_name=current_secret_name, - new_secret_name=new_secret_name, - new_secret_value=new_secret_value, - ) - - # Verify error response - assert response["status"] == "error" - assert "mismatch" in response["message"].lower() - assert "expected-value" in response["message"] diff --git a/tests/litellm_utils_tests/test_secret_manager.py b/tests/litellm_utils_tests/test_secret_manager.py deleted file mode 100644 index 6d088b0c0e2..00000000000 --- a/tests/litellm_utils_tests/test_secret_manager.py +++ /dev/null @@ -1,103 +0,0 @@ -import hashlib -import json -import os - -from dotenv import load_dotenv - -load_dotenv() -import tempfile -from typing import Final - -import pytest - -import litellm -from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 -from litellm.secret_managers.main import ( - get_secret, -) - -_AWS_FIXTURE_MASTER_KEY_SHA256: Final = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" - - -def load_vertex_ai_credentials(): - # Define the path to the vertex_key.json file - print("loading vertex ai credentials") - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - # Create a temporary file - with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: - # Write the updated content to the temporary files - json.dump(service_account_key_data, temp_file, indent=2) - - # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) - - -def test_aws_secret_manager(): - import json - - AWSSecretsManagerV2.load_aws_secret_manager(use_aws_secret_manager=True) - - secret_val = get_secret("litellm_master_key") - - print(f"secret_val: {secret_val}") - - # cast json to dict - secret_val = json.loads(secret_val) - - assert ( - hashlib.sha256(secret_val["litellm_master_key"].encode()).hexdigest() == _AWS_FIXTURE_MASTER_KEY_SHA256 - ), "Expected the fixture value stored in the CI AWS account's litellm_master_key secret" - - -def redact_oidc_signature(secret_val): - # remove the last part of `.` and replace it with "SIGNATURE_REMOVED" - return secret_val.split(".")[:-1] + ["SIGNATURE_REMOVED"] - - -@pytest.mark.skipif( - os.environ.get("K_SERVICE") is None, - reason="Cannot run without being in GCP Cloud Run", -) -def test_oidc_google(): - secret_val = get_secret( - "oidc/google/https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" - ) - - print(f"secret_val: {redact_oidc_signature(secret_val)}") -@pytest.mark.skipif( - os.environ.get("ACTIONS_ID_TOKEN_REQUEST_TOKEN") is None, - reason="Cannot run without being in GitHub Actions", -) -def test_oidc_github(): - secret_val = get_secret( - "oidc/github/https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" - ) - - print(f"secret_val: {redact_oidc_signature(secret_val)}") diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py deleted file mode 100644 index 5ec7b96dcfa..00000000000 --- a/tests/litellm_utils_tests/test_utils.py +++ /dev/null @@ -1,81 +0,0 @@ -import copy -import logging -from unittest import mock - -from dotenv import load_dotenv - -from litellm.types.utils import StandardCallbackDynamicParams - -load_dotenv() -import os - -import pytest - -import litellm -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, headers -from litellm.utils import ( - get_llm_provider, - get_supported_openai_params, - get_token_count, - trim_messages, - validate_environment, -) -from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - -# Assuming your trim_messages, shorten_message_to_fit_limit, and get_token_count functions are all in a module named 'message_utils' -@pytest.fixture(autouse=True) -def reset_mock_cache(): - from litellm.utils import _model_cache - - _model_cache.flush_cache() - - -# Test 1: Check trimming of normal message - - -# test_basic_trimming() - - -# test_basic_trimming_no_max_tokens_specified() - - -# test_multiple_messages_trimming() - - -# test_multiple_messages_no_trimming() - - -# test_large_trimming() - - -# test_get_valid_models() - - -# test validate environment - - -from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.utils import get_applied_guardrails -from unittest.mock import Mock - - -def test_get_whitelisted_models(): - """ - Snapshot of all bedrock models as of 12/24/2024. - - Enforce any new bedrock chat model to be added as `bedrock_converse` unless explicitly whitelisted. - - Create whitelist to prevent naming regressions for older litellm versions. - """ - whitelisted_models = [] - for model, info in litellm.model_cost.items(): - if info.get("litellm_provider") == "bedrock" and info.get("mode") == "chat": - whitelisted_models.append(model) - - # Write to a local file - with open("whitelisted_bedrock_models.txt", "w") as file: - for model in whitelisted_models: - file.write(f"{model}\n") - - print("whitelisted_models written to whitelisted_bedrock_models.txt") diff --git a/tests/litellm_utils_tests/vertex_key.json b/tests/litellm_utils_tests/vertex_key.json deleted file mode 100644 index 800969fb305..00000000000 --- a/tests/litellm_utils_tests/vertex_key.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "type": "service_account", - "project_id": "litellm-ci-cd", - "private_key_id": "", - "private_key": "", - "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", - "client_id": "116563532503305622785", - "auth_uri": "https://accounts.google.com/o/oauth2/auth", - "token_uri": "https://oauth2.googleapis.com/token", - "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", - "universe_domain": "googleapis.com" -} diff --git a/tests/unified_google_tests/base_google_genai_proxy_sdk_test.py b/tests/unified_google_tests/base_google_genai_proxy_sdk_test.py deleted file mode 100644 index 50d419c704b..00000000000 --- a/tests/unified_google_tests/base_google_genai_proxy_sdk_test.py +++ /dev/null @@ -1,140 +0,0 @@ -from __future__ import annotations - -import os -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional - -import pytest - -from tests._master_key import MASTER_KEY - -try: - from google import genai - from google.genai import types - - GOOGLE_GENAI_SDK_AVAILABLE = True -except ImportError: - GOOGLE_GENAI_SDK_AVAILABLE = False - -PROMPT = "Reply with only the single word: pong" - - -def has_vertex_credentials() -> bool: - credentials_file = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "") - if credentials_file and os.path.isfile(credentials_file): - return True - return bool( - os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - and os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - ) - - -def _make_client(proxy_url: str) -> "genai.Client": - return genai.Client( - api_key=MASTER_KEY, - http_options={"base_url": proxy_url}, - ) - - -def _generation_config() -> "types.GenerateContentConfig": - return types.GenerateContentConfig( - temperature=0, - top_p=0.95, - top_k=20, - ) - - -def _collect_stream_text(chunks: List["types.GenerateContentResponse"]) -> str: - return "".join(chunk.text for chunk in chunks if chunk.text) - - -class BaseGoogleGenAIProxySDKTest(ABC): - @property - @abstractmethod - def proxy_model_name(self) -> str: ... - - @property - @abstractmethod - def model_config(self) -> Dict[str, Any]: ... - - def _skip_reason_if_credentials_missing(self) -> Optional[str]: - model = self.model_config.get("model", "") - if model.startswith("gemini/"): - if not os.getenv("GEMINI_API_KEY"): - return "GEMINI_API_KEY not set — skipping Gemini proxy SDK tests" - return None - - if "vertex_ai" in model: - if has_vertex_credentials(): - return None - return "Vertex AI credentials not set — skipping Vertex AI proxy SDK tests" - - return f"Unsupported model for proxy SDK tests: {model}" - - def _require_proxy_sdk(self) -> None: - if not GOOGLE_GENAI_SDK_AVAILABLE: - pytest.skip("google-genai SDK not installed") - reason = self._skip_reason_if_credentials_missing() - if reason: - pytest.skip(reason) - - def test_proxy_genai_sdk_non_streaming(self, google_genai_proxy_url: str) -> None: - self._require_proxy_sdk() - - client = _make_client(google_genai_proxy_url) - response = client.models.generate_content( - model=self.proxy_model_name, - contents=types.Part.from_text(text=PROMPT), - config=_generation_config(), - ) - - assert response is not None - assert response.text is not None - assert len(response.text.strip()) > 0 - - def test_proxy_genai_sdk_streaming_completes_without_errors( - self, google_genai_proxy_url: str - ) -> None: - self._require_proxy_sdk() - - client = _make_client(google_genai_proxy_url) - stream = client.models.generate_content_stream( - model=self.proxy_model_name, - contents=types.Part.from_text(text=PROMPT), - config=_generation_config(), - ) - - chunks: List[types.GenerateContentResponse] = [] - stream_error: Optional[Exception] = None - - try: - for chunk in stream: - chunks.append(chunk) - except Exception as exc: - stream_error = exc - - assert ( - stream_error is None - ), f"Streaming raised {type(stream_error).__name__}: {stream_error}" - assert len(chunks) > 0, "Expected at least one streaming chunk" - assert _collect_stream_text(chunks).strip(), "Expected non-empty streamed text" - - def test_proxy_genai_sdk_streaming_dict_style( - self, google_genai_proxy_url: str - ) -> None: - self._require_proxy_sdk() - - client = _make_client(google_genai_proxy_url) - stream = client.models.generate_content_stream( - model=self.proxy_model_name, - contents={"text": PROMPT}, - config={ - "temperature": 0, - "top_p": 0.95, - "top_k": 20, - }, - ) - - chunks = list(stream) - assert len(chunks) > 0 - assert _collect_stream_text(chunks).strip() diff --git a/tests/unified_google_tests/base_google_test.py b/tests/unified_google_tests/base_google_test.py deleted file mode 100644 index 558ec2445fb..00000000000 --- a/tests/unified_google_tests/base_google_test.py +++ /dev/null @@ -1,276 +0,0 @@ -import asyncio -import json -import os -import tempfile -from typing import Any, AsyncIterator, Dict, List, Optional, Union -import pytest - - -import litellm -from litellm.google_genai import ( - generate_content, - agenerate_content, - agenerate_content_stream, -) -from google.genai.types import ContentDict, PartDict -from litellm.types.google_genai.main import GenerateContentResponse -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingPayload - - -def load_vertex_ai_credentials(model: str): - """Load Vertex AI credentials for tests""" - # Define the path to the vertex_key.json file - if "vertex_ai" not in model: - return None - print("loading vertex ai credentials") - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - # Create a temporary file - with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: - # Write the updated content to the temporary files - json.dump(service_account_key_data, temp_file, indent=2) - - # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) - - return os.path.abspath(temp_file.name) - - -class TestCustomLogger(CustomLogger): - def __init__( - self, - ): - self.standard_logging_object: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - print("in async_log_success_event") - print("kwargs=", json.dumps(kwargs, indent=4, default=str)) - self.standard_logging_object = kwargs["standard_logging_object"] - pass - - -class BaseGoogleGenAITest: - """Base class for Google GenAI generate content tests to reduce code duplication""" - - @property - def model_config(self) -> Dict[str, Any]: - """Override in subclasses to provide model-specific configuration""" - raise NotImplementedError("Subclasses must implement model_config") - - @property - def _temp_files_to_cleanup(self): - """Lazy initialization of temp files list""" - if not hasattr(self, "_temp_files_list"): - self._temp_files_list = [] - return self._temp_files_list - - def cleanup_temp_files(self): - """Clean up any temporary files created during testing""" - for temp_file in self._temp_files_to_cleanup: - try: - os.unlink(temp_file) - except OSError: - pass # File might already be deleted - self._temp_files_to_cleanup.clear() - - def _validate_non_streaming_response(self, response: Any): - """Validate non-streaming response structure""" - # Handle type checking - response should be a GenerateContentResponse for non-streaming - if isinstance(response, AsyncIterator): - pytest.fail("Expected non-streaming response but got AsyncIterator") - - assert isinstance( - response, GenerateContentResponse - ), f"Expected GenerateContentResponse, got {type(response)}" - print(f"Response: {response.model_dump_json(indent=4)}") - - # Basic validation - adjust based on actual Google GenAI response structure - # The exact structure may vary, so we'll be flexible here - assert response is not None, "Response should not be None" - - def _validate_streaming_response(self, chunks: List[Any]): - """Validate streaming response chunks""" - assert isinstance(chunks, list), f"Expected list of chunks, got {type(chunks)}" - assert len(chunks) >= 0, "Should have at least 0 chunks" - print(f"Total chunks received: {len(chunks)}") - - def _validate_standard_logging_payload( - self, slp: StandardLoggingPayload, response: Any - ): - """ - Validate that a StandardLoggingPayload object matches the expected response for Google GenAI - - Args: - slp (StandardLoggingPayload): The standard logging payload object to validate - response: The Google GenAI response to compare against - """ - # Validate payload exists - assert slp is not None, "Standard logging payload should not be None" - - # Validate basic structure - assert ( - "prompt_tokens" in slp - ), "Standard logging payload should have prompt_tokens" - assert ( - "completion_tokens" in slp - ), "Standard logging payload should have completion_tokens" - assert ( - "total_tokens" in slp - ), "Standard logging payload should have total_tokens" - assert ( - "response_cost" in slp - ), "Standard logging payload should have response_cost" - - # Validate token counts are reasonable (non-negative numbers) - assert slp["prompt_tokens"] >= 0, "Prompt tokens should be non-negative" - assert slp["completion_tokens"] >= 0, "Completion tokens should be non-negative" - assert slp["total_tokens"] >= 0, "Total tokens should be non-negative" - - # Validate spend - assert slp["response_cost"] >= 0, "Response cost should be non-negative" - - print( - f"Standard logging payload validation passed: prompt_tokens={slp['prompt_tokens']}, completion_tokens={slp['completion_tokens']}, total_tokens={slp['total_tokens']}, cost={slp['response_cost']}" - ) - - @pytest.mark.parametrize("is_async", [False, True]) - @pytest.mark.asyncio - async def test_non_streaming_base(self, is_async: bool): - """Base test for non-streaming requests (parametrized for sync/async)""" - request_params = self.model_config - contents = ContentDict( - parts=[PartDict(text="Hello, can you tell me a short joke?")], - role="user", - ) - temp_file_path = load_vertex_ai_credentials(model=request_params["model"]) - if temp_file_path: - self._temp_files_to_cleanup.append(temp_file_path) - - litellm.turn_on_debug() - - print( - f"Testing {'async' if is_async else 'sync'} non-streaming with model config: {request_params}" - ) - print(f"Contents: {contents}") - - if is_async: - print("\n--- Testing async agenerate_content ---") - response = await agenerate_content(contents=contents, **request_params) - else: - print("\n--- Testing sync generate_content ---") - response = generate_content(contents=contents, **request_params) - - print( - f"{'Async' if is_async else 'Sync'} response: {json.dumps(response, indent=2, default=str)}" - ) - self._validate_non_streaming_response(response) - - return response - - @pytest.mark.asyncio - async def test_async_non_streaming_with_logging(self): - """Test async non-streaming Google GenAI generate content with logging""" - litellm.turn_on_debug() - litellm.logging_callback_manager._reset_all_callbacks() - litellm.set_verbose = True - test_custom_logger = TestCustomLogger() - litellm.callbacks = [test_custom_logger] - - request_params = self.model_config - temp_file_path = load_vertex_ai_credentials(model=request_params["model"]) - if temp_file_path: - self._temp_files_to_cleanup.append(temp_file_path) - contents = ContentDict( - parts=[PartDict(text="Hello, can you tell me a short joke?")], - role="user", - ) - - print("\n--- Testing async agenerate_content with logging ---") - response = await agenerate_content(contents=contents, **request_params) - - print("Google GenAI response=", json.dumps(response, indent=4, default=str)) - - print("sleeping for 5 seconds...") - await asyncio.sleep(5) - print( - "standard logging payload=", - json.dumps( - test_custom_logger.standard_logging_object, indent=4, default=str - ), - ) - - assert response is not None - assert test_custom_logger.standard_logging_object is not None - - self._validate_standard_logging_payload( - test_custom_logger.standard_logging_object, response - ) - - @pytest.mark.asyncio - async def test_async_streaming_with_logging(self): - """Test async streaming Google GenAI generate content with logging""" - litellm.turn_on_debug() - litellm.set_verbose = True - litellm.logging_callback_manager._reset_all_callbacks() - test_custom_logger = TestCustomLogger() - litellm.callbacks = [test_custom_logger] - - request_params = self.model_config - temp_file_path = load_vertex_ai_credentials(model=request_params["model"]) - if temp_file_path: - self._temp_files_to_cleanup.append(temp_file_path) - contents = ContentDict( - parts=[PartDict(text="Hello, can you tell me a short joke?")], - role="user", - ) - - print("\n--- Testing async agenerate_content_stream with logging ---") - response = await agenerate_content_stream(contents=contents, **request_params) - - chunks = [] - async for chunk in response: - print(f"Google GenAI chunk: {chunk}") - chunks.append(chunk) - - print("sleeping for 5 seconds...") - await asyncio.sleep(5) - print( - "standard logging payload=", - json.dumps( - test_custom_logger.standard_logging_object, indent=4, default=str - ), - ) - - assert len(chunks) >= 0 - assert test_custom_logger.standard_logging_object is not None - - self._validate_standard_logging_payload( - test_custom_logger.standard_logging_object, chunks - ) diff --git a/tests/unified_google_tests/base_interactions_test.py b/tests/unified_google_tests/base_interactions_test.py deleted file mode 100644 index 4dff91e697e..00000000000 --- a/tests/unified_google_tests/base_interactions_test.py +++ /dev/null @@ -1,117 +0,0 @@ -""" -Abstract base class for Interactions API tests. - -This class provides common test cases that can be inherited by provider-specific -test classes. Subclasses must implement get_model() and get_api_key(). -""" - -import os -from abc import ABC, abstractmethod - -import pytest -import litellm -import litellm.interactions as interactions - - -class BaseInteractionsTest(ABC): - """Abstract base class for interactions API tests. - - Subclasses must implement get_model() and get_api_key(). - All test methods are inherited and run against the specific provider. - """ - - @abstractmethod - def get_model(self) -> str: - """Return the model string for this provider.""" - pass - - @abstractmethod - def get_api_key(self) -> str: - """Return the API key for this provider.""" - pass - - def test_create_simple_string_input(self): - """Test creating an interaction with a simple string input.""" - litellm.turn_on_debug() - api_key = self.get_api_key() - if not api_key: - pytest.skip(f"API key not set for {self.__class__.__name__}") - - response = interactions.create( - model=self.get_model(), - input="Hello, what is 2 + 2?", - api_key=api_key, - ) - assert response is not None - assert response.id is not None or response.status is not None - - # Check outputs per OpenAPI spec - if response.outputs: - assert len(response.outputs) > 0 - - # Check usage per OpenAPI spec - # The spec defines: total_input_tokens, total_output_tokens - if response.usage: - # Usage is a dict in InteractionsAPIResponse - if isinstance(response.usage, dict): - assert ( - response.usage.get("total_input_tokens") is not None - or response.usage.get("total_output_tokens") is not None - ) - else: - # If it's an object, check attributes - assert hasattr(response.usage, "total_input_tokens") or hasattr( - response.usage, "total_output_tokens" - ) - - def test_create_with_system_instruction(self): - """Test creating an interaction with system_instruction.""" - api_key = self.get_api_key() - if not api_key: - pytest.skip(f"API key not set for {self.__class__.__name__}") - - response = interactions.create( - model=self.get_model(), - input="What are you?", - system_instruction="You are a helpful pirate assistant. Always respond like a pirate.", - api_key=api_key, - ) - assert response is not None - # Verify the response reflects the system instruction - if response.outputs: - assert len(response.outputs) > 0 - - def test_create_streaming(self): - """Test creating a streaming interaction.""" - api_key = self.get_api_key() - if not api_key: - pytest.skip(f"API key not set for {self.__class__.__name__}") - - response_stream = interactions.create( - model=self.get_model(), - input="Count from 1 to 3.", - stream=True, - api_key=api_key, - ) - - # Collect all chunks - chunks = [] - for chunk in response_stream: - chunks.append(chunk) - - assert len(chunks) > 0 - - @pytest.mark.asyncio - async def test_acreate_simple(self): - """Test async interaction creation.""" - api_key = self.get_api_key() - if not api_key: - pytest.skip(f"API key not set for {self.__class__.__name__}") - - response = await interactions.acreate( - model=self.get_model(), - input="What is the speed of light?", - api_key=api_key, - ) - assert response is not None - assert response.id is not None or response.status is not None diff --git a/tests/unified_google_tests/conftest.py b/tests/unified_google_tests/conftest.py deleted file mode 100644 index ad0e135ec04..00000000000 --- a/tests/unified_google_tests/conftest.py +++ /dev/null @@ -1,221 +0,0 @@ -# conftest.py - -import asyncio -import importlib -import os -import socket -import threading -import time -from pathlib import Path -from typing import Final, Iterator, Tuple - -import pytest -import uvicorn -from dotenv import load_dotenv - -from tests._master_key import MASTER_KEY - -load_dotenv() - -import litellm # noqa: E402,F401 - -from tests._vcr_conftest_common import ( # noqa: E402,F401 - VerboseReporterState, - _pin_multipart_boundary, - apply_vcr_auto_marker_to_items, - emit_cassette_cache_session_banner, - emit_vcr_classification_summary, - emit_vcr_diagnostic_log, - install_live_call_probe, - record_vcr_outcome, - register_persister_if_enabled, - reset_vcr_diag_dir, - vcr_config_dict, -) - -_verbose_state = VerboseReporterState() - -PROXY_CONFIG_PATH = Path(__file__).parent / "google_genai_proxy_test_config.yaml" -PROXY_MASTER_KEY: Final = MASTER_KEY -PROXY_START_TIMEOUT_S = 30.0 - - -def _start_proxy_server( - config_path: str, -) -> Tuple[str, uvicorn.Server, threading.Thread, socket.socket]: - from litellm.proxy.proxy_server import ( - app as proxy_app, - cleanup_router_config_variables, - initialize, - ) - - cleanup_router_config_variables() - - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - sock.bind(("127.0.0.1", 0)) - host, port = sock.getsockname() - - config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning") - server = uvicorn.Server(config) - - def _run() -> None: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - loop.run_until_complete(initialize(config=config_path, debug=True)) - loop.run_until_complete(server.serve(sockets=[sock])) - - thread = threading.Thread(target=_run, daemon=True) - thread.start() - - start_time = time.time() - while not server.started: - if not thread.is_alive(): - raise RuntimeError("LiteLLM proxy failed to start") - if time.time() - start_time > PROXY_START_TIMEOUT_S: - raise TimeoutError("LiteLLM proxy did not start in time") - time.sleep(0.05) - - return f"http://{host}:{port}", server, thread, sock - - -@pytest.fixture(scope="session") -def google_genai_proxy_url() -> Iterator[str]: - from base_google_genai_proxy_sdk_test import has_vertex_credentials - from base_google_test import load_vertex_ai_credentials - - saved_env = { - key: os.environ.get(key) - for key in ( - "DATABASE_URL", - "DIRECT_URL", - "LITELLM_MASTER_KEY", - "STORE_MODEL_IN_DB", - "GOOGLE_APPLICATION_CREDENTIALS", - ) - } - temp_credentials_path: str | None = None - os.environ.pop("DATABASE_URL", None) - os.environ.pop("DIRECT_URL", None) - os.environ["LITELLM_MASTER_KEY"] = PROXY_MASTER_KEY - os.environ["STORE_MODEL_IN_DB"] = "False" - - if has_vertex_credentials(): - credentials_file = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "") - if not (credentials_file and os.path.isfile(credentials_file)): - vertex_credentials_path = load_vertex_ai_credentials( - model="vertex_ai/gemini-3.5-flash-lite" - ) - if vertex_credentials_path: - temp_credentials_path = vertex_credentials_path - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = vertex_credentials_path - - server_url, server, thread, sock = _start_proxy_server(str(PROXY_CONFIG_PATH)) - try: - yield server_url - finally: - server.should_exit = True - thread.join(timeout=10) - sock.close() - if temp_credentials_path: - try: - os.unlink(temp_credentials_path) - except OSError: - pass - for key, value in saved_env.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - - -@pytest.fixture(scope="session") -def event_loop(): - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - yield loop - loop.close() - - -@pytest.fixture(scope="function", autouse=True) -def setup_and_teardown(request): - """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - - - if "google_genai_proxy_url" not in request.fixturenames: - importlib.reload(litellm) - - loop = asyncio.get_event_loop_policy().new_event_loop() - asyncio.set_event_loop(loop) - print(litellm) - yield - - # Teardown code (executes after the yield point) - loop.close() # Close the loop created earlier - asyncio.set_event_loop(None) # Remove the reference to the loop - - -@pytest.fixture(scope="module") -def vcr_config(): - return vcr_config_dict() - - -def pytest_recording_configure(config, vcr): - register_persister_if_enabled(vcr) - - -@pytest.hookimpl(hookwrapper=True) -def pytest_runtest_makereport(item, call): - outcome = yield - rep = outcome.get_result() - setattr(item, f"rep_{rep.when}", rep) - - -@pytest.fixture(autouse=True) -def _vcr_outcome_gate(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) - - -def pytest_configure(config): - _verbose_state.remember_pluginmanager(config) - reset_vcr_diag_dir() - - -def pytest_runtest_logreport(report): - _verbose_state.maybe_emit_verdict(report) - - -def pytest_collection_modifyitems(config, items): - apply_vcr_auto_marker_to_items( - items, - skip_nodeid_suffixes=( - "test_proxy_genai_sdk_non_streaming", - "test_proxy_genai_sdk_streaming_completes_without_errors", - "test_proxy_genai_sdk_streaming_dict_style", - ), - ) - - # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests - custom_logger_tests = [ - item for item in items if "custom_logger" in item.parent.name - ] - other_tests = [item for item in items if "custom_logger" not in item.parent.name] - - # Sort tests based on their names - custom_logger_tests.sort(key=lambda x: x.name) - other_tests.sort(key=lambda x: x.name) - - # Reorder the items list - items[:] = custom_logger_tests + other_tests - - -def pytest_terminal_summary(terminalreporter, exitstatus, config): - emit_cassette_cache_session_banner(terminalreporter) - emit_vcr_classification_summary(terminalreporter) - emit_vcr_diagnostic_log(terminalreporter) diff --git a/tests/unified_google_tests/google_genai_proxy_test_config.yaml b/tests/unified_google_tests/google_genai_proxy_test_config.yaml deleted file mode 100644 index 268ab2028f8..00000000000 --- a/tests/unified_google_tests/google_genai_proxy_test_config.yaml +++ /dev/null @@ -1,21 +0,0 @@ -model_list: - - model_name: gemini-3.5-flash-lite - litellm_params: - model: gemini/gemini-3.5-flash-lite - api_key: os.environ/GEMINI_API_KEY - - - model_name: vertex-gemini-3.5-flash-lite - litellm_params: - model: vertex_ai/gemini-3.5-flash-lite - vertex_location: global - -router_settings: - retry_policy: - RateLimitErrorRetries: 5 - -general_settings: - master_key: os.environ/LITELLM_MASTER_KEY - store_model_in_db: false - -litellm_settings: - drop_params: true diff --git a/tests/unified_google_tests/test_gemini_interactions.py b/tests/unified_google_tests/test_gemini_interactions.py deleted file mode 100644 index 0e0719843f1..00000000000 --- a/tests/unified_google_tests/test_gemini_interactions.py +++ /dev/null @@ -1,23 +0,0 @@ -""" -Tests for Gemini Interactions API. - -Inherits from BaseInteractionsTest to run the same test suite against Gemini. -""" - -import os - -from tests.unified_google_tests.base_interactions_test import ( - BaseInteractionsTest, -) - - -class TestGeminiInteractions(BaseInteractionsTest): - """Test Gemini Interactions API using the base test suite.""" - - def get_model(self) -> str: - """Return the Gemini model string.""" - return "gemini/gemini-2.5-flash" - - def get_api_key(self) -> str: - """Return the Gemini API key from environment.""" - return os.getenv("GEMINI_API_KEY", "") diff --git a/tests/unified_google_tests/test_google_ai_studio.py b/tests/unified_google_tests/test_google_ai_studio.py deleted file mode 100644 index 0d9f8365ecf..00000000000 --- a/tests/unified_google_tests/test_google_ai_studio.py +++ /dev/null @@ -1,473 +0,0 @@ -from base_google_genai_proxy_sdk_test import BaseGoogleGenAIProxySDKTest -from base_google_test import BaseGoogleGenAITest - -import pytest -import litellm -import unittest.mock -import json - - -class TestGoogleGenAIStudio(BaseGoogleGenAITest, BaseGoogleGenAIProxySDKTest): - """Test Google GenAI Studio""" - - test_non_streaming_base = None - - @property - def model_config(self): - return { - "model": "gemini/gemini-3.5-flash-lite", - } - - @property - def proxy_model_name(self) -> str: - return "gemini-3.5-flash-lite" - - -@pytest.mark.asyncio -async def test_mock_stream_generate_content_with_tools(): - """Test streaming function call response parsing and validation""" - from litellm.types.google_genai.main import ToolConfigDict - - litellm.turn_on_debug() - contents = [ - { - "role": "user", - "parts": [ - { - "text": "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning" - } - ], - } - ] - - # Mock streaming response chunks that represent a function call response - mock_response_chunk = { - "candidates": [ - { - "content": { - "parts": [ - { - "functionCall": { - "name": "schedule_meeting", - "args": { - "attendees": ["Bob", "Alice"], - "date": "2025-03-27", - "time": "10:00", - "topic": "Q3 planning", - }, - } - } - ], - "role": "model", - }, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 15, - "candidatesTokenCount": 5, - "totalTokenCount": 20, - }, - } - - # Mock the HTTP handler - with unittest.mock.patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=unittest.mock.AsyncMock, - ) as mock_post: - # Create mock response object - mock_response = unittest.mock.MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - - # Mock aiter_lines: yield one line at a time (no trailing newlines), - # with a blank line between events, matching httpx aiter_lines behaviour. - async def mock_aiter_lines(): - yield f"data: {json.dumps(mock_response_chunk)}" - yield "" - yield "data: [DONE]" - yield "" - - mock_response.aiter_lines = mock_aiter_lines - mock_post.return_value = mock_response - - print( - "\n--- Testing async agenerate_content_stream with function call parsing ---" - ) - response = await litellm.google_genai.agenerate_content_stream( - model="gemini/gemini-3.5-flash-lite", - contents=contents, - tools=[ - { - "functionDeclarations": [ - { - "name": "schedule_meeting", - "description": "Schedules a meeting with specified attendees at a given time and date.", - "parameters": { - "type": "object", - "properties": { - "attendees": { - "type": "array", - "items": {"type": "string"}, - "description": "List of people attending the meeting.", - }, - "date": { - "type": "string", - "description": "Date of the meeting (e.g., '2024-07-29')", - }, - "time": { - "type": "string", - "description": "Time of the meeting (e.g., '15:00')", - }, - "topic": { - "type": "string", - "description": "The subject or topic of the meeting.", - }, - }, - "required": ["attendees", "date", "time", "topic"], - }, - } - ] - } - ], - ) - - # Collect all chunks and parse function calls - chunks = [] - function_calls = [] - - chunk_count = 0 - async for chunk in response: - chunk_count += 1 - print(f"Received chunk {chunk_count}: {chunk}") - chunks.append(chunk) - - # Stop after a reasonable number of chunks to prevent infinite loop - if chunk_count > 10: - break - - # Parse function calls from byte chunks - if isinstance(chunk, bytes): - try: - # Decode bytes to string - chunk_str = chunk.decode("utf-8") - print(f"Decoded chunk: {chunk_str}") - - # Extract JSON from Server-Sent Events format (data: {...}) - if chunk_str.startswith("data: ") and not chunk_str.startswith( - "data: [DONE]" - ): - json_str = chunk_str[6:].strip() # Remove 'data: ' prefix - try: - parsed_json = json.loads(json_str) - print(f"Parsed JSON: {parsed_json}") - - # Parse function calls from the JSON - if "candidates" in parsed_json: - for candidate in parsed_json["candidates"]: - if ( - "content" in candidate - and "parts" in candidate["content"] - ): - for part in candidate["content"]["parts"]: - if "functionCall" in part: - function_calls.append( - { - "name": part["functionCall"][ - "name" - ], - "args": part["functionCall"][ - "args" - ], - } - ) - print( - f"Found function call: {part['functionCall']}" - ) - except json.JSONDecodeError as e: - print(f"Failed to parse JSON: {e}") - except UnicodeDecodeError as e: - print(f"Failed to decode bytes: {e}") - - # Handle dict responses (in case some chunks are already parsed) - elif isinstance(chunk, dict): - # Direct dict response - if "candidates" in chunk: - for candidate in chunk["candidates"]: - if "content" in candidate and "parts" in candidate["content"]: - for part in candidate["content"]["parts"]: - if "functionCall" in part: - function_calls.append( - { - "name": part["functionCall"]["name"], - "args": part["functionCall"]["args"], - } - ) - - # Handle object responses with attributes - elif hasattr(chunk, "candidates") and chunk.candidates: - for candidate in chunk.candidates: - if hasattr(candidate, "content") and candidate.content: - if ( - hasattr(candidate.content, "parts") - and candidate.content.parts - ): - for part in candidate.content.parts: - if ( - hasattr(part, "function_call") - and part.function_call - ): - function_calls.append( - { - "name": part.function_call.name, - "args": part.function_call.args, - } - ) - - # Assertions - print(f"\nFunction calls found: {function_calls}") - print(f"Total chunks received: {chunk_count}") - - # Assert we found at least one function call - assert ( - len(function_calls) > 0 - ), "Expected at least one function call in the streaming response" - - # Check the first function call - function_call = function_calls[0] - - # Assert function name - assert ( - function_call["name"] == "schedule_meeting" - ), f"Expected function name 'schedule_meeting', got '{function_call['name']}'" - - # Assert function arguments - args = function_call["args"] - assert "attendees" in args, "Expected 'attendees' in function call arguments" - assert "date" in args, "Expected 'date' in function call arguments" - assert "time" in args, "Expected 'time' in function call arguments" - assert "topic" in args, "Expected 'topic' in function call arguments" - - # Assert specific argument values - assert args["attendees"] == [ - "Bob", - "Alice", - ], f"Expected attendees ['Bob', 'Alice'], got {args['attendees']}" - assert ( - args["date"] == "2025-03-27" - ), f"Expected date '2025-03-27', got {args['date']}" - assert args["time"] == "10:00", f"Expected time '10:00', got {args['time']}" - assert ( - args["topic"] == "Q3 planning" - ), f"Expected topic 'Q3 planning', got {args['topic']}" - - print("✅ All function call assertions passed!") - - -@pytest.mark.asyncio -async def test_validate_post_request_parameters(): - """ - Test that the correct parameters are sent in the POST request to Google GenAI API - - Params validated - 1. model - 2. contents - 3. tools - """ - from litellm.types.google_genai.main import ToolConfigDict - - contents = [ - { - "role": "user", - "parts": [ - { - "text": "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning" - } - ], - } - ] - - tools = [ - { - "functionDeclarations": [ - { - "name": "schedule_meeting", - "description": "Schedules a meeting with specified attendees at a given time and date.", - "parameters": { - "type": "object", - "properties": { - "attendees": { - "type": "array", - "items": {"type": "string"}, - "description": "List of people attending the meeting.", - }, - "date": { - "type": "string", - "description": "Date of the meeting (e.g., '2024-07-29')", - }, - "time": { - "type": "string", - "description": "Time of the meeting (e.g., '15:00')", - }, - "topic": { - "type": "string", - "description": "The subject or topic of the meeting.", - }, - }, - "required": ["attendees", "date", "time", "topic"], - }, - } - ] - } - ] - - # Mock the HTTP handler to capture the request - with unittest.mock.patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=unittest.mock.AsyncMock, - ) as mock_post: - # Create mock response object - mock_response = unittest.mock.MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - - # Mock aiter_lines: yield one line at a time (no trailing newlines), - # with a blank line between events, matching httpx aiter_lines behaviour. - async def mock_aiter_lines(): - yield "data: [DONE]" - yield "" - - mock_response.aiter_lines = mock_aiter_lines - mock_post.return_value = mock_response - - print("\n--- Testing POST request parameters validation ---") - - # Make the API call - response = await litellm.google_genai.agenerate_content_stream( - model="gemini/gemini-3.5-flash-lite", contents=contents, tools=tools - ) - - # Consume the response to ensure the request is made - async for chunk in response: - pass - - # Validate that the HTTP post was called - assert mock_post.called, "Expected HTTP POST to be called" - - # Get the call arguments - call_args, call_kwargs = mock_post.call_args - - print(f"POST call args: {call_args}") - print(f"POST call kwargs: {call_kwargs}") - - # Validate URL contains the correct endpoint - if call_args: - url = call_args[0] if len(call_args) > 0 else call_kwargs.get("url") - assert url is not None, "Expected URL to be provided" - assert ( - "generativelanguage.googleapis.com" in url - ), f"Expected Google API URL, got: {url}" - assert ( - "streamGenerateContent" in url - ), f"Expected streamGenerateContent endpoint, got: {url}" - print(f"✅ URL validation passed: {url}") - - # Get the request data/json from the call - request_data = None - if "data" in call_kwargs: - # If data is passed as bytes, decode it - if isinstance(call_kwargs["data"], bytes): - request_data = json.loads(call_kwargs["data"].decode("utf-8")) - else: - request_data = call_kwargs["data"] - elif "json" in call_kwargs: - request_data = call_kwargs["json"] - - assert request_data is not None, "Expected request data to be provided" - print(f"Request data: {json.dumps(request_data, indent=2)}") - - # Validate model field - assert "model" in request_data, "Expected 'model' field in request data" - # Model might be transformed, but should contain gemini-3.5-flash-lite - model_value = request_data["model"] - assert ( - "gemini-3.5-flash-lite" in model_value - ), f"Expected model to contain 'gemini-3.5-flash-lite', got: {model_value}" - print(f"✅ Model validation passed: {model_value}") - - # Validate contents field - assert "contents" in request_data, "Expected 'contents' field in request data" - request_contents = request_data["contents"] - assert isinstance(request_contents, list), "Expected contents to be a list" - assert len(request_contents) > 0, "Expected at least one content item" - - # Check the first content item - first_content = request_contents[0] - assert "role" in first_content, "Expected 'role' in content item" - assert ( - first_content["role"] == "user" - ), f"Expected role 'user', got: {first_content['role']}" - assert "parts" in first_content, "Expected 'parts' in content item" - assert isinstance(first_content["parts"], list), "Expected parts to be a list" - assert len(first_content["parts"]) > 0, "Expected at least one part" - - # Check the text content - first_part = first_content["parts"][0] - assert "text" in first_part, "Expected 'text' in part" - expected_text = "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning" - assert ( - first_part["text"] == expected_text - ), f"Expected text '{expected_text}', got: {first_part['text']}" - print(f"✅ Contents validation passed") - - # Validate tools field - assert "tools" in request_data, "Expected 'tools' field in request data" - request_tools = request_data["tools"] - assert isinstance(request_tools, list), "Expected tools to be a list" - assert len(request_tools) > 0, "Expected at least one tool" - - # Check the first tool - first_tool = request_tools[0] - assert ( - "functionDeclarations" in first_tool - ), "Expected 'functionDeclarations' in tool" - function_declarations = first_tool["functionDeclarations"] - assert isinstance( - function_declarations, list - ), "Expected functionDeclarations to be a list" - assert ( - len(function_declarations) > 0 - ), "Expected at least one function declaration" - - # Check the function declaration - func_decl = function_declarations[0] - assert "name" in func_decl, "Expected 'name' in function declaration" - assert ( - func_decl["name"] == "schedule_meeting" - ), f"Expected function name 'schedule_meeting', got: {func_decl['name']}" - assert ( - "description" in func_decl - ), "Expected 'description' in function declaration" - assert ( - "parameters" in func_decl - ), "Expected 'parameters' in function declaration" - - # Check function parameters - params = func_decl["parameters"] - assert "type" in params, "Expected 'type' in parameters" - assert ( - params["type"] == "object" - ), f"Expected parameters type 'object', got: {params['type']}" - assert "properties" in params, "Expected 'properties' in parameters" - assert "required" in params, "Expected 'required' in parameters" - - # Check required fields - required_fields = params["required"] - expected_required = ["attendees", "date", "time", "topic"] - assert set(required_fields) == set( - expected_required - ), f"Expected required fields {expected_required}, got: {required_fields}" - print(f"✅ Tools validation passed") - - print("✅ All POST request parameter validations passed!") diff --git a/tests/unified_google_tests/test_litellm_responses_bridge.py b/tests/unified_google_tests/test_litellm_responses_bridge.py deleted file mode 100644 index d32e0cccc73..00000000000 --- a/tests/unified_google_tests/test_litellm_responses_bridge.py +++ /dev/null @@ -1,30 +0,0 @@ -""" -Tests for LiteLLM Responses bridge provider. - -Inherits from BaseInteractionsTest to run the same test suite against -the litellm_responses bridge provider, which calls litellm.responses() internally. -""" - -import os - -from tests.unified_google_tests.base_interactions_test import ( - BaseInteractionsTest, -) - - -class TestLiteLLMResponsesBridge(BaseInteractionsTest): - """Test LiteLLM Responses bridge using the base test suite.""" - - test_create_streaming = None - - def get_model(self) -> str: - """Return the model string for the bridge provider. - - The bridge provider uses litellm.responses() internally, so we can - use any model that litellm.responses() supports (e.g., gpt-5.5). - """ - return "gpt-5.5" - - def get_api_key(self) -> str: - """Return the OpenAI API key from environment.""" - return os.getenv("OPENAI_API_KEY", "") diff --git a/tests/unified_google_tests/test_vertex_ai_native.py b/tests/unified_google_tests/test_vertex_ai_native.py deleted file mode 100644 index 4bdc1a55283..00000000000 --- a/tests/unified_google_tests/test_vertex_ai_native.py +++ /dev/null @@ -1,17 +0,0 @@ -from base_google_genai_proxy_sdk_test import BaseGoogleGenAIProxySDKTest -from base_google_test import BaseGoogleGenAITest - - -class TestVertexAIGenerateContent(BaseGoogleGenAITest, BaseGoogleGenAIProxySDKTest): - """Test Vertex AI""" - - @property - def model_config(self): - return { - "model": "vertex_ai/gemini-3.5-flash-lite", - "vertex_location": "global", - } - - @property - def proxy_model_name(self) -> str: - return "vertex-gemini-3.5-flash-lite" diff --git a/tests/unified_google_tests/test_vertex_anthropic.py b/tests/unified_google_tests/test_vertex_anthropic.py deleted file mode 100644 index f11ee28aacb..00000000000 --- a/tests/unified_google_tests/test_vertex_anthropic.py +++ /dev/null @@ -1,369 +0,0 @@ -import asyncio -import json -from typing import Any, AsyncIterator, Dict, List, Optional, Union -import pytest -from unittest.mock import MagicMock, AsyncMock, patch -import httpx - - -import litellm -from litellm.google_genai import agenerate_content, agenerate_content_stream -from google.genai.types import ContentDict, PartDict, GenerateContentResponse -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingPayload - - -async def vertex_anthropic_mock_response(*args, **kwargs): - """Mock response for vertex AI anthropic call""" - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"Content-Type": "application/json"} - mock_response.json.return_value = { - "id": "msg_vrtx_013Wki5RFQXAspL7rmxRFjZg", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-4", - "content": [ - { - "type": "text", - "text": "Why don't scientists trust atoms? Because they make up everything!", - } - ], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 15, "output_tokens": 20}, - } - return mock_response - - -@pytest.mark.asyncio -async def test_vertex_anthropic_mocked(): - """Test agenerate_content with mocked HTTP calls to validate URL and request body""" - - # Set up test data - contents = ContentDict( - parts=[PartDict(text="Hello, can you tell me a short joke?")], - role="user", - ) - - # Expected values for validation - expected_url = "https://us-east5-aiplatform.googleapis.com/v1/projects/internal-litellm-local-dev/locations/us-east5/publishers/anthropic/models/claude-sonnet-4:rawPredict" - expected_body_keys = {"messages", "anthropic_version", "max_tokens"} - expected_message_content = "Hello, can you tell me a short joke?" - - # Patch the AsyncHTTPHandler.post method at the module level - with patch( - "litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = await vertex_anthropic_mock_response() - - response = await agenerate_content( - contents=contents, - model="vertex_ai/claude-sonnet-4", - vertex_location="us-east5", - vertex_project="internal-litellm-local-dev", - custom_llm_provider="vertex_ai", - ) - - # Verify the call was made - assert mock_post.call_count == 1 - - # Get the call arguments - call_args = mock_post.call_args - call_kwargs = call_args.kwargs if call_args else {} - - # Extract URL (could be in args[0] or kwargs['url']) - if call_args and len(call_args[0]) > 0: - actual_url = call_args[0][0] - else: - actual_url = call_kwargs.get("url", "") - - # Validate URL - print(f"Expected URL: {expected_url}") - print(f"Actual URL: {actual_url}") - assert ( - actual_url == expected_url - ), f"Expected URL {expected_url}, but got {actual_url}" - - # Validate headers - actual_headers = call_kwargs.get("headers", {}) - print(f"Actual headers: {actual_headers}") - - # Validate Authorization header exists - auth_header_found = any( - k.lower() == "authorization" for k in actual_headers.keys() - ) - assert ( - auth_header_found - ), f"Authorization header should be present. Found headers: {list(actual_headers.keys())}" - - # Validate request body - request_body = None - if "data" in call_kwargs: - request_body = ( - json.loads(call_kwargs["data"]) - if isinstance(call_kwargs["data"], str) - else call_kwargs["data"] - ) - elif "json" in call_kwargs: - request_body = call_kwargs["json"] - - print(f"Request body: {json.dumps(request_body, indent=2)}") - assert request_body is not None, "Request body should not be None" - - # Validate required keys in request body - actual_body_keys = set(request_body.keys()) - assert expected_body_keys.issubset( - actual_body_keys - ), f"Expected keys {expected_body_keys} not found in {actual_body_keys}" - - # Validate message content - messages = request_body.get("messages", []) - assert len(messages) > 0, "Messages should not be empty" - assert ( - messages[0]["role"] == "user" - ), f"Expected first message role to be 'user', got {messages[0]['role']}" - - # Check message content structure - content = messages[0]["content"] - if isinstance(content, list): - text_content = next( - (item["text"] for item in content if item.get("type") == "text"), None - ) - else: - text_content = content - - assert ( - text_content == expected_message_content - ), f"Expected message content '{expected_message_content}', got '{text_content}'" - - # Validate anthropic_version - assert ( - request_body["anthropic_version"] == "vertex-2023-10-16" - ), f"Expected anthropic_version 'vertex-2023-10-16', got {request_body['anthropic_version']}" - - # Validate max_tokens - assert ( - "max_tokens" in request_body - ), "max_tokens should be present in request body" - assert isinstance( - request_body["max_tokens"], int - ), f"max_tokens should be integer, got {type(request_body['max_tokens'])}" - - print("✅ All validations passed!") - print(f"Response: {response}") - - -class MockAsyncStreamResponse: - """Mock async streaming response that mimics httpx streaming response""" - - def __init__(self): - self.status_code = 200 - self.headers = {"Content-Type": "text/event-stream"} - self._chunks = [ - { - "type": "message_start", - "message": { - "id": "msg_vrtx_013Wki5RFQXAspL7rmxRFjZg", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-4", - "content": [], - "stop_reason": None, - "stop_sequence": None, - "usage": {"input_tokens": 15, "output_tokens": 0}, - }, - }, - { - "type": "content_block_start", - "index": 0, - "content_block": {"type": "text", "text": ""}, - }, - { - "type": "content_block_delta", - "index": 0, - "delta": { - "type": "text_delta", - "text": "Why don't scientists trust atoms? ", - }, - }, - { - "type": "content_block_delta", - "index": 0, - "delta": { - "type": "text_delta", - "text": "Because they make up everything!", - }, - }, - {"type": "content_block_stop", "index": 0}, - { - "type": "message_delta", - "delta": {"stop_reason": "end_turn", "stop_sequence": None}, - "usage": {"output_tokens": 20}, - }, - {"type": "message_stop"}, - ] - - async def aiter_bytes(self, chunk_size=1024): - """Async iterator for response bytes""" - for chunk in self._chunks: - yield f"data: {json.dumps(chunk)}\n\n".encode() - - async def aiter_lines(self): - """Async iterator for response lines (required by anthropic handler)""" - for chunk in self._chunks: - yield f"data: {json.dumps(chunk)}\n\n" - - -async def vertex_anthropic_streaming_mock_response(*args, **kwargs): - """Mock streaming response for vertex AI anthropic call""" - return MockAsyncStreamResponse() - - -@pytest.mark.asyncio -async def test_vertex_anthropic_streaming_mocked(): - """Test agenerate_content_stream with mocked HTTP calls to validate URL and request body""" - - # Set up test data - contents = ContentDict( - parts=[PartDict(text="Hello, can you tell me a short joke?")], - role="user", - ) - - # Expected values for validation (same as non-streaming) - expected_url = "https://us-east5-aiplatform.googleapis.com/v1/projects/internal-litellm-local-dev/locations/us-east5/publishers/anthropic/models/claude-sonnet-4:streamRawPredict" - expected_body_keys = {"messages", "anthropic_version", "max_tokens"} - expected_message_content = "Hello, can you tell me a short joke?" - - # Patch the AsyncHTTPHandler.post method at the module level - with patch( - "litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = await vertex_anthropic_streaming_mock_response() - - response_stream = await agenerate_content_stream( - contents=contents, - model="vertex_ai/claude-sonnet-4", - vertex_location="us-east5", - vertex_project="internal-litellm-local-dev", - custom_llm_provider="vertex_ai", - ) - - # Verify the call was made - assert mock_post.call_count == 1 - - # Get the call arguments - call_args = mock_post.call_args - call_kwargs = call_args.kwargs if call_args else {} - - # Extract URL (could be in args[0] or kwargs['url']) - if call_args and len(call_args[0]) > 0: - actual_url = call_args[0][0] - else: - actual_url = call_kwargs.get("url", "") - - # Validate URL (same as non-streaming) - print(f"Expected URL: {expected_url}") - print(f"Actual URL: {actual_url}") - assert ( - actual_url == expected_url - ), f"Expected URL {expected_url}, but got {actual_url}" - - # Validate headers - actual_headers = call_kwargs.get("headers", {}) - print(f"Actual headers: {actual_headers}") - - # Validate Authorization header exists - auth_header_found = any( - k.lower() == "authorization" for k in actual_headers.keys() - ) - assert ( - auth_header_found - ), f"Authorization header should be present. Found headers: {list(actual_headers.keys())}" - - # Validate anthropic-version header exists and has correct value - anthropic_version_found = False - for header_name, header_value in actual_headers.items(): - if header_name.lower() == "anthropic-version": - assert ( - header_value == "2023-06-01" - ), f"Expected anthropic-version: 2023-06-01, but got {header_value}" - anthropic_version_found = True - break - assert anthropic_version_found, "anthropic-version header should be present" - - # Validate content-type and accept headers - content_type_found = any( - k.lower() == "content-type" for k in actual_headers.keys() - ) - accept_found = any(k.lower() == "accept" for k in actual_headers.keys()) - assert content_type_found, "content-type header should be present" - assert accept_found, "accept header should be present" - - # Validate request body (same structure as non-streaming) - request_body = None - if "data" in call_kwargs: - request_body = ( - json.loads(call_kwargs["data"]) - if isinstance(call_kwargs["data"], str) - else call_kwargs["data"] - ) - elif "json" in call_kwargs: - request_body = call_kwargs["json"] - - print(f"Request body: {json.dumps(request_body, indent=2)}") - assert request_body is not None, "Request body should not be None" - - # Validate required keys in request body - actual_body_keys = set(request_body.keys()) - assert expected_body_keys.issubset( - actual_body_keys - ), f"Expected keys {expected_body_keys} not found in {actual_body_keys}" - - # Validate message content - messages = request_body.get("messages", []) - assert len(messages) > 0, "Messages should not be empty" - assert ( - messages[0]["role"] == "user" - ), f"Expected first message role to be 'user', got {messages[0]['role']}" - - # Check message content structure - content = messages[0]["content"] - if isinstance(content, list): - text_content = next( - (item["text"] for item in content if item.get("type") == "text"), None - ) - else: - text_content = content - - assert ( - text_content == expected_message_content - ), f"Expected message content '{expected_message_content}', got '{text_content}'" - - # Validate anthropic_version in body - assert ( - request_body["anthropic_version"] == "vertex-2023-10-16" - ), f"Expected anthropic_version 'vertex-2023-10-16', got {request_body['anthropic_version']}" - - # Validate max_tokens - assert ( - "max_tokens" in request_body - ), "max_tokens should be present in request body" - assert isinstance( - request_body["max_tokens"], int - ), f"max_tokens should be integer, got {type(request_body['max_tokens'])}" - - # Test that we can iterate over the streaming response - chunks_received = [] - try: - async for chunk in response_stream: - chunks_received.append(chunk) - print(f"Received streaming chunk: {chunk}") - except Exception as e: - print(f"Note: Streaming iteration might not work with mock response: {e}") - - print(f"✅ All streaming validations passed!") - print(f"Total chunks received: {len(chunks_received)}") - print(f"Response stream: {response_stream}") diff --git a/tests/unified_google_tests/vertex_key.json b/tests/unified_google_tests/vertex_key.json deleted file mode 100644 index 800969fb305..00000000000 --- a/tests/unified_google_tests/vertex_key.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "type": "service_account", - "project_id": "litellm-ci-cd", - "private_key_id": "", - "private_key": "", - "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", - "client_id": "116563532503305622785", - "auth_uri": "https://accounts.google.com/o/oauth2/auth", - "token_uri": "https://oauth2.googleapis.com/token", - "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", - "universe_domain": "googleapis.com" -} diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index 626d53d3384..bf5e4d31f52 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -1,13 +1,17 @@ import base64 import json +from collections.abc import Mapping from typing import Final import httpx import pytest +from pydantic import JsonValue, TypeAdapter from respx import MockRouter import litellm -import litellm.interactions as interactions +from litellm import interactions + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @pytest.fixture @@ -92,3 +96,163 @@ class TestInteractionsAcreateOffline: assert "What is the speed of light?" in serialized assert "response_id:resp-offline" in base64.b64decode(response.id.removeprefix("resp_")).decode() assert response.status == "completed" + + +@pytest.mark.parametrize( + ("model", "api_key", "url", "response_body", "expected_body"), + [ + ( + "gemini/gemini-2.5-flash", + "gemini-offline", + "https://generativelanguage.googleapis.com/v1beta/interactions", + { + "id": "interaction-sync", + "object": "interaction", + "model": "gemini-2.5-flash", + "status": "completed", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": "four"}]}], + "usage": {"input_tokens": 5, "output_tokens": 1}, + }, + {"model": "gemini-2.5-flash", "input": "What is 2 + 2?"}, + ), + ( + "gpt-4o", + "sk-offline", + "https://api.openai.com/v1/responses", + { + "id": "resp-sync", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "four"}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6}, + }, + {"model": "gpt-4o", "input": "What is 2 + 2?"}, + ), + ], +) +def test_create_simple_string_input_for_google_and_responses( + respx_mock: MockRouter, + model: str, + api_key: str, + url: str, + response_body: Mapping[str, JsonValue], + expected_body: Mapping[str, JsonValue], +) -> None: + route: Final = respx_mock.post(url).mock(return_value=httpx.Response(200, json=response_body)) + + response: Final = interactions.create( + model=model, + input="What is 2 + 2?", + api_key=api_key, + ) + + assert _JSON_OBJECT.validate_json(route.calls.last.request.content) == expected_body + assert response.status == "completed" + + +@pytest.mark.parametrize( + ("model", "api_key", "url", "response_body", "expected_body"), + [ + ( + "gemini/gemini-2.5-flash", + "gemini-offline", + "https://generativelanguage.googleapis.com/v1beta/interactions", + { + "id": "interaction-instruction", + "object": "interaction", + "model": "gemini-2.5-flash", + "status": "completed", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": "ahoy"}]}], + }, + { + "model": "gemini-2.5-flash", + "input": "What are you?", + "system_instruction": "Answer like a pirate.", + }, + ), + ( + "gpt-4o", + "sk-offline", + "https://api.openai.com/v1/responses", + { + "id": "resp-instruction", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "ahoy"}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6}, + }, + {"model": "gpt-4o", "input": "What are you?", "instructions": "Answer like a pirate."}, + ), + ], +) +def test_create_with_system_instruction_for_google_and_responses( + respx_mock: MockRouter, + model: str, + api_key: str, + url: str, + response_body: Mapping[str, JsonValue], + expected_body: Mapping[str, JsonValue], +) -> None: + route: Final = respx_mock.post(url).mock(return_value=httpx.Response(200, json=response_body)) + + response: Final = interactions.create( + model=model, + input="What are you?", + system_instruction="Answer like a pirate.", + api_key=api_key, + ) + + assert _JSON_OBJECT.validate_json(route.calls.last.request.content) == expected_body + assert response.status == "completed" + + +def test_create_streaming_returns_interaction_events( + respx_mock: MockRouter, +) -> None: + route: Final = respx_mock.post( + "https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse" + ).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"event_type":"step.delta","delta":{"type":"text_delta","text":"1, 2, 3"}}\n\n' + b'data: {"event_type":"interaction.completed","id":"interaction-stream","status":"completed"}\n\n' + ), + ) + ) + + chunks: Final = tuple( + interactions.create( + model="gemini/gemini-2.5-flash", + input="Count from 1 to 3.", + stream=True, + api_key="gemini-offline", + ) + ) + + assert tuple((chunk.event_type, chunk.delta) for chunk in chunks) == ( + ("step.delta", {"type": "text_delta", "text": "1, 2, 3"}), + ("interaction.completed", None), + ) + assert _JSON_OBJECT.validate_json(route.calls.last.request.content) == { + "model": "gemini-2.5-flash", + "input": "Count from 1 to 3.", + "stream": True, + } diff --git a/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py b/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py new file mode 100644 index 00000000000..b4c654ee175 --- /dev/null +++ b/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py @@ -0,0 +1,241 @@ +from types import SimpleNamespace +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue, TypeAdapter +from respx import MockRouter + +import litellm +from litellm.google_genai import agenerate_content_stream +from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PROMPT: Final = "Hello, can you tell me a short joke?" +_CONTENTS: Final = [{"role": "user", "parts": [{"text": _PROMPT}]}] +_FUNCTION_ARGS: Final = { + "attendees": ["Bob", "Alice"], + "date": "2025-03-27", + "time": "10:00", + "topic": "Q3 planning", +} + + +@pytest.mark.parametrize("model", ("gemini-2.5-flash", "vertex_ai/gemini-2.5-flash")) +def test_generate_content_request_maps_provider_parameters(model: str) -> None: + request: Final = GoogleGenAIConfig().transform_generate_content_request( + model=model, + contents=_CONTENTS, + tools=None, + generate_content_config_dict={"temperature": 0.4, "topP": 0.8, "topK": 12}, + ) + + assert _JSON_OBJECT.validate_python(request) == { + "model": model, + "contents": _CONTENTS, + "tools": None, + "generationConfig": {"temperature": 0.4, "topP": 0.8, "topK": 12}, + } + + +@pytest.mark.parametrize("model", ("gemini-2.5-flash", "vertex_ai/gemini-2.5-flash")) +def test_generate_content_response_keeps_text_usage_and_logging_response(model: str) -> None: + raw_response: Final = httpx.Response( + 200, + json={ + "responseId": "scripted-response", + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "Scientists trust atoms."}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 15, "candidatesTokenCount": 5, "totalTokenCount": 20}, + }, + ) + logging_obj: Final = SimpleNamespace(model_call_details={}) + response: Final = GoogleGenAIConfig().transform_generate_content_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert _JSON_OBJECT.validate_python(response.candidates[0]) == { + "content": {"role": "model", "parts": [{"text": "Scientists trust atoms."}]}, + "finishReason": "STOP", + } + assert _JSON_OBJECT.validate_python(response.model_dump(by_alias=True))["usageMetadata"] == { + "promptTokenCount": 15, + "candidatesTokenCount": 5, + "totalTokenCount": 20, + } + assert logging_obj.model_call_details == {"httpx_response": raw_response} + + +def test_generate_content_response_parses_function_call() -> None: + raw_response: Final = httpx.Response( + 200, + json={ + "responseId": "scripted-tool-call", + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "schedule_meeting", "args": _FUNCTION_ARGS}}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 15, "candidatesTokenCount": 5, "totalTokenCount": 20}, + }, + ) + logging_obj: Final = SimpleNamespace(model_call_details={}) + response: Final = GoogleGenAIConfig().transform_generate_content_response( + model="gemini-2.5-flash", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert _JSON_OBJECT.validate_python(response.candidates[0]) == { + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "schedule_meeting", "args": _FUNCTION_ARGS}}], + }, + "finishReason": "STOP", + } + + +def test_generate_content_request_preserves_tool_declarations() -> None: + tools: Final = [ + { + "functionDeclarations": [ + { + "name": "schedule_meeting", + "description": "Schedules a meeting.", + "parameters": { + "type": "object", + "properties": {"topic": {"type": "string"}}, + "required": ["topic"], + }, + } + ] + } + ] + request: Final = GoogleGenAIConfig().transform_generate_content_request( + model="gemini-2.5-flash", + contents=_CONTENTS, + tools=tools, + generate_content_config_dict={"temperature": 0.4}, + ) + + assert _JSON_OBJECT.validate_python(request) == { + "model": "gemini-2.5-flash", + "contents": _CONTENTS, + "tools": tools, + "generationConfig": {"temperature": 0.4}, + } + + +@pytest.mark.asyncio +async def test_generate_content_stream_returns_text_and_usage( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse" + ).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"responseId":"scripted-stream","candidates":[{"content":{"role":"model","parts":' + b'[{"text":"Scientists trust "}]}}]}\n\n' + b'data: {"responseId":"scripted-stream","candidates":[{"content":{"role":"model","parts":' + b'[{"text":"atoms."}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":15,' + b'"candidatesTokenCount":5,"totalTokenCount":20}}\n\n' + ), + ) + ) + + chunks: Final = tuple( + [ + chunk + async for chunk in await agenerate_content_stream( + model="gemini/gemini-2.5-flash", + contents=_CONTENTS, + api_key="gemini-offline", + ) + ] + ) + frames: Final = tuple(_JSON_OBJECT.validate_json(chunk.removeprefix(b"data: ").strip()) for chunk in chunks) + + assert frames == ( + { + "responseId": "scripted-stream", + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "Scientists trust "}]}} + ], + }, + { + "responseId": "scripted-stream", + "candidates": [{"content": {"role": "model", "parts": [{"text": "atoms."}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 15, "candidatesTokenCount": 5, "totalTokenCount": 20}, + }, + ) + assert _JSON_OBJECT.validate_json(route.calls.last.request.content) == { + "model": "gemini-2.5-flash", + "contents": _CONTENTS, + "tools": None, + "generationConfig": {}, + } + + +@pytest.mark.asyncio +async def test_generate_content_stream_parses_function_call( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse" + ).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"responseId":"scripted-tool-call","candidates":[{"content":{"role":"model","parts":' + b'[{"functionCall":{"name":"schedule_meeting","args":{"attendees":["Bob","Alice"],' + b'"date":"2025-03-27","time":"10:00","topic":"Q3 planning"}}}]}}],"usageMetadata":' + b'{"promptTokenCount":15,"candidatesTokenCount":5,"totalTokenCount":20}}\n\n' + ), + ) + ) + tools: Final = [{"functionDeclarations": [{"name": "schedule_meeting", "parameters": {"type": "object"}}]}] + + chunks: Final = tuple( + [ + chunk + async for chunk in await agenerate_content_stream( + model="gemini/gemini-2.5-flash", + contents=_CONTENTS, + tools=tools, + api_key="gemini-offline", + ) + ] + ) + frames: Final = tuple(_JSON_OBJECT.validate_json(chunk.removeprefix(b"data: ").strip()) for chunk in chunks) + + assert frames == ( + { + "responseId": "scripted-tool-call", + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"functionCall": {"name": "schedule_meeting", "args": _FUNCTION_ARGS}}], + } + } + ], + "usageMetadata": {"promptTokenCount": 15, "candidatesTokenCount": 5, "totalTokenCount": 20}, + }, + ) + assert _JSON_OBJECT.validate_json(route.calls.last.request.content)["tools"] == tools diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 0208d87b354..0153717f910 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -12,6 +12,8 @@ from litellm.anthropic_beta_headers_manager import ( from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.vertex_ai import VertexPartnerProvider @pytest.mark.parametrize( @@ -31,6 +33,44 @@ def test_vertex_ai_anthropic_thinking_param(model, expected_thinking): assert "thinking" not in supported_openai_params +@pytest.mark.parametrize( + ("stream", "operation"), + [(False, "rawPredict"), (True, "streamRawPredict")], +) +def test_vertex_anthropic_request_body_and_endpoint(stream: bool, operation: str) -> None: + headers: Final[dict[str, str]] = {} + body: Final = VertexAIAnthropicConfig().transform_request( + model="claude-sonnet-4", + messages=[{"role": "user", "content": "Hello, can you tell me a short joke?"}], + optional_params={"max_tokens": 100, "anthropic_version": "vertex-2023-10-16"}, + litellm_params={}, + headers=headers, + ) + url: Final = VertexBase.create_vertex_url( + vertex_location="us-east5", + vertex_project="scripted-project", + partner=VertexPartnerProvider.claude, + stream=stream, + model="claude-sonnet-4", + ) + + assert body == { + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "Hello, can you tell me a short joke?"}], + } + ], + "max_tokens": 100, + "anthropic_version": "vertex-2023-10-16", + } + assert headers == {} + assert url == ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/scripted-project/locations/us-east5/" + f"publishers/anthropic/models/claude-sonnet-4:{operation}" + ) + + def test_get_supported_params_thinking(): config = VertexAIAnthropicConfig() params = config.get_supported_openai_params(model="claude-sonnet-4") diff --git a/tests/unit/secret_managers/test_hashicorp_secret_manager.py b/tests/unit/secret_managers/test_hashicorp_secret_manager.py index fc18cb8b8f7..21facc749f8 100644 --- a/tests/unit/secret_managers/test_hashicorp_secret_manager.py +++ b/tests/unit/secret_managers/test_hashicorp_secret_manager.py @@ -4,12 +4,14 @@ from collections.abc import Mapping from pathlib import Path from typing import Final +import httpx import pytest import respx from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization -from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.hazmat.primitives.asymmetric import ec from cryptography.x509.oid import NameOID +from pydantic import JsonValue, TypeAdapter import litellm.proxy.proxy_server from litellm.secret_managers.hashicorp_secret_manager import HashicorpSecretManager @@ -17,6 +19,7 @@ from litellm.secret_managers.hashicorp_secret_manager import HashicorpSecretMana VAULT_ADDR: Final = "http://vault.test:8200" LOGIN_RESPONSE: Final = {"auth": {"client_token": "hvs.login-token", "lease_duration": 3600}} SECRET_RESPONSE: Final = {"data": {"data": {"key": "sk-from-vault", "password": "pw-from-vault"}}} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) NAMESPACE_ENV_VARS: Final = ("HCP_VAULT_NAMESPACE", "HCP_VAULT_LOGIN_NAMESPACE", "HCP_VAULT_SECRET_NAMESPACE") PARITY_ENV_VARS: Final = ( @@ -42,7 +45,10 @@ def _build_manager(monkeypatch: pytest.MonkeyPatch, env: Mapping[str, str]) -> H monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) for name in NAMESPACE_ENV_VARS: monkeypatch.delenv(name, raising=False) + for name in ("HCP_VAULT_CLIENT_CERT", "HCP_VAULT_CLIENT_KEY", "HCP_VAULT_CERT_ROLE"): + monkeypatch.delenv(name, raising=False) monkeypatch.setenv("HCP_VAULT_ADDR", VAULT_ADDR) + monkeypatch.delenv("HCP_VAULT_TOKEN", raising=False) monkeypatch.setenv("HCP_VAULT_APPROLE_ROLE_ID", "role-id") monkeypatch.setenv("HCP_VAULT_APPROLE_SECRET_ID", "secret-id") for name, value in env.items(): @@ -210,17 +216,18 @@ async def test_async_write_and_read_share_the_secret_namespace_target(monkeypatc def _write_self_signed_cert(directory: Path) -> tuple[Path, Path]: - private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + private_key: Final = ec.derive_private_key(1, ec.SECP256R1()) name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "litellm-test")]) - now: Final = datetime.datetime.now(datetime.timezone.utc) + valid_from: Final = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc) + valid_until: Final = datetime.datetime(2049, 12, 31, tzinfo=datetime.timezone.utc) certificate: Final = ( x509.CertificateBuilder() .subject_name(name) .issuer_name(name) .public_key(private_key.public_key()) - .serial_number(x509.random_serial_number()) - .not_valid_before(now) - .not_valid_after(now + datetime.timedelta(days=1)) + .serial_number(1) + .not_valid_before(valid_from) + .not_valid_after(valid_until) .sign(private_key, hashes.SHA256()) ) cert_path: Final = directory / "client.crt" @@ -242,11 +249,13 @@ def test_tls_login_uses_login_namespace(monkeypatch: pytest.MonkeyPatch, tmp_pat monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) for name in NAMESPACE_ENV_VARS: monkeypatch.delenv(name, raising=False) + monkeypatch.delenv("HCP_VAULT_TOKEN", raising=False) monkeypatch.delenv("HCP_VAULT_APPROLE_ROLE_ID", raising=False) monkeypatch.delenv("HCP_VAULT_APPROLE_SECRET_ID", raising=False) monkeypatch.setenv("HCP_VAULT_ADDR", VAULT_ADDR) monkeypatch.setenv("HCP_VAULT_CLIENT_CERT", str(cert)) monkeypatch.setenv("HCP_VAULT_CLIENT_KEY", str(key)) + monkeypatch.setenv("HCP_VAULT_CERT_ROLE", "test-role") monkeypatch.setenv("HCP_VAULT_NAMESPACE", "admin") monkeypatch.setenv("HCP_VAULT_LOGIN_NAMESPACE", "root") manager: Final = HashicorpSecretManager() @@ -254,6 +263,349 @@ def test_tls_login_uses_login_namespace(monkeypatch: pytest.MonkeyPatch, tmp_pat assert manager._auth_via_tls_cert() == "hvs.login-token" assert login_route.calls.last.request.headers["X-Vault-Namespace"] == "root" + assert _JSON_OBJECT.validate_json(login_route.calls.last.request.content) == {"name": "test-role"} + assert manager.cache.get_cache("hcp_vault_token") == "hvs.login-token" + + +@respx.mock +def test_hashicorp_secret_manager_get_secret(monkeypatch: pytest.MonkeyPatch) -> None: + manager: Final = _build_manager(monkeypatch, {"HCP_VAULT_NAMESPACE": "admin"}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + read_route: Final = respx.get(f"{VAULT_ADDR}/v1/admin/secret/data/sample-secret-mock").respond( + json=SECRET_RESPONSE + ) + + assert manager.sync_read_secret("sample-secret-mock") == "sk-from-vault" + + assert read_route.call_count == 1 + assert read_route.calls.last.request.headers["X-Vault-Token"] == "hvs.login-token" + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_write_secret(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {"HCP_VAULT_NAMESPACE": "admin"}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + write_route: Final = respx.post(f"{VAULT_ADDR}/v1/admin/secret/data/sample-secret").respond( + json={"data": {"version": 1}} + ) + + response: Final = await manager.async_write_secret("sample-secret", "value-mock") + + assert response == {"data": {"version": 1}} + assert write_route.call_count == 1 + assert _JSON_OBJECT.validate_json(write_route.calls.last.request.content) == {"data": {"key": "value-mock"}} + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_write_secret_with_team_overrides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + write_route: Final = respx.post( + f"{VAULT_ADDR}/v1/team-namespace/kv-team/data/teams/custom/team-secret" + ).respond(json={"data": {"version": 1}}) + settings: Final = { + "secret_manager_settings": { + "namespace": "team-namespace", + "mount": "kv-team", + "path_prefix": "teams/custom", + "data": "password", + } + } + + response: Final = await manager.async_write_secret( + "team-secret", "value-mock", optional_params=settings + ) + + assert response == {"data": {"version": 1}} + assert _JSON_OBJECT.validate_json(write_route.calls.last.request.content) == {"data": {"password": "value-mock"}} + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_delete_secret(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {"HCP_VAULT_NAMESPACE": "admin"}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + delete_route: Final = respx.delete(f"{VAULT_ADDR}/v1/admin/secret/data/sample-secret").respond( + status_code=204 + ) + + response: Final = await manager.async_delete_secret("sample-secret") + + assert response == { + "status": "success", + "message": "Secret sample-secret deleted successfully", + } + assert delete_route.call_count == 1 + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_delete_secret_with_team_overrides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + delete_route: Final = respx.delete( + f"{VAULT_ADDR}/v1/team-namespace/kv-team/data/teams/custom/team-secret" + ).respond(status_code=204) + settings: Final = { + "secret_manager_settings": { + "namespace": "team-namespace", + "mount": "kv-team", + "path_prefix": "teams/custom", + } + } + + response: Final = await manager.async_delete_secret("team-secret", optional_params=settings) + + assert response == { + "status": "success", + "message": "Secret team-secret deleted successfully", + } + assert delete_route.call_count == 1 + + +@respx.mock +def test_hashicorp_secret_manager_approle_auth(monkeypatch: pytest.MonkeyPatch) -> None: + manager: Final = _build_manager( + monkeypatch, + { + "HCP_VAULT_APPROLE_ROLE_ID": "test-role-id-123", + "HCP_VAULT_APPROLE_SECRET_ID": "test-secret-id-456", + "HCP_VAULT_APPROLE_MOUNT_PATH": "approle", + }, + ) + login_route: Final = respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + + assert manager._auth_via_approle() == "hvs.login-token" + + assert _JSON_OBJECT.validate_json(login_route.calls.last.request.content) == { + "role_id": "test-role-id-123", + "secret_id": "test-secret-id-456", + } + assert manager.cache.get_cache("hcp_vault_approle_token") == "hvs.login-token" + + +def test_hashicorp_custom_mount_and_prefix(monkeypatch: pytest.MonkeyPatch) -> None: + manager: Final = _build_manager(monkeypatch, {"HCP_VAULT_NAMESPACE": "admin"}) + + assert manager.get_url("my-secret") == f"{VAULT_ADDR}/v1/admin/secret/data/my-secret" + assert manager.get_url("my-secret", mount_name="kv") == f"{VAULT_ADDR}/v1/admin/kv/data/my-secret" + assert manager.get_url("my-secret", path_prefix="myapp") == ( + f"{VAULT_ADDR}/v1/admin/secret/data/myapp/my-secret" + ) + assert manager.get_url("my-secret", mount_name="kv", path_prefix="production") == ( + f"{VAULT_ADDR}/v1/admin/kv/data/production/my-secret" + ) + + +@pytest.mark.parametrize( + "malicious_secret_name", + [ + "../../../other-app/creds", + "litellm/../../secret", + "foo\nbar", + "foo\u2028bar", + "foo\u2029bar", + "foo\x85bar", + ], +) +def test_hashicorp_get_url_rejects_path_traversal( + monkeypatch: pytest.MonkeyPatch, malicious_secret_name: str +) -> None: + manager: Final = _build_manager(monkeypatch, {}) + + with pytest.raises(ValueError, match="Invalid secret_name"): + manager.get_url(malicious_secret_name) + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_rotate_secret_different_names( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + current_route: Final = respx.get(f"{VAULT_ADDR}/v1/secret/data/old-secret").respond( + json={"data": {"data": {"key": "old-secret-value"}}} + ) + write_route: Final = respx.post(f"{VAULT_ADDR}/v1/secret/data/new-secret").respond( + json={"data": {"version": 1}} + ) + new_route: Final = respx.get(f"{VAULT_ADDR}/v1/secret/data/new-secret").respond( + json={"data": {"data": {"key": "new-secret-value"}}} + ) + delete_route: Final = respx.delete(f"{VAULT_ADDR}/v1/secret/data/old-secret").respond( + status_code=204 + ) + + response: Final = await manager.async_rotate_secret( + "old-secret", "new-secret", "new-secret-value" + ) + + assert response == {"data": {"version": 1}} + assert current_route.call_count == 1 + assert _JSON_OBJECT.validate_json(write_route.calls.last.request.content) == { + "data": {"key": "new-secret-value", "description": "Rotated from old-secret"} + } + assert new_route.call_count == 1 + assert delete_route.call_count == 1 + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_rotate_secret_same_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + secret_route: Final = respx.get(f"{VAULT_ADDR}/v1/secret/data/same-secret").mock( + side_effect=[ + httpx.Response(200, json={"data": {"data": {"key": "old-secret-value"}}}), + httpx.Response(200, json={"data": {"data": {"key": "updated-secret-value"}}}), + ] + ) + write_route: Final = respx.post(f"{VAULT_ADDR}/v1/secret/data/same-secret").respond( + json={"data": {"version": 1}} + ) + delete_route: Final = respx.delete(f"{VAULT_ADDR}/v1/secret/data/same-secret").respond( + status_code=204 + ) + + response: Final = await manager.async_rotate_secret( + "same-secret", "same-secret", "updated-secret-value" + ) + + assert response == {"data": {"version": 1}} + assert secret_route.call_count == 2 + assert write_route.call_count == 1 + assert delete_route.call_count == 0 + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_rotate_secret_current_not_found( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + respx.get(f"{VAULT_ADDR}/v1/secret/data/non-existent").respond( + status_code=404, text="Not Found" + ) + + response: Final = await manager.async_rotate_secret( + "non-existent", "new-secret", "new-value" + ) + + assert response["status"] == "error" + assert "non-existent" in response["message"] + assert "not found" in response["message"].lower() + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_rotate_secret_write_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + respx.get(f"{VAULT_ADDR}/v1/secret/data/old-secret").respond( + json={"data": {"data": {"key": "old-secret-value"}}} + ) + write_route: Final = respx.post(f"{VAULT_ADDR}/v1/secret/data/new-secret").respond( + json={"status": "error", "message": "Write failed"} + ) + + response: Final = await manager.async_rotate_secret( + "old-secret", "new-secret", "new-value" + ) + + assert response == {"status": "error", "message": "Write failed"} + assert write_route.call_count == 1 + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_rotate_secret_with_team_overrides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + current_route: Final = respx.get( + f"{VAULT_ADDR}/v1/team-namespace/kv-team/data/teams/custom/team-old-secret" + ).respond(json={"data": {"data": {"password": "old-secret-value"}}}) + write_route: Final = respx.post( + f"{VAULT_ADDR}/v1/team-namespace/kv-team/data/teams/custom/team-new-secret" + ).respond(json={"data": {"version": 1}}) + new_route: Final = respx.get( + f"{VAULT_ADDR}/v1/team-namespace/kv-team/data/teams/custom/team-new-secret" + ).respond(json={"data": {"data": {"password": "new-team-secret-value"}}}) + delete_route: Final = respx.delete( + f"{VAULT_ADDR}/v1/team-namespace/kv-team/data/teams/custom/team-old-secret" + ).respond(status_code=204) + settings: Final = { + "secret_manager_settings": { + "namespace": "team-namespace", + "mount": "kv-team", + "path_prefix": "teams/custom", + "data": "password", + } + } + + response: Final = await manager.async_rotate_secret( + "team-old-secret", "team-new-secret", "new-team-secret-value", optional_params=settings + ) + + assert response == {"data": {"version": 1}} + assert current_route.call_count == 1 + assert _JSON_OBJECT.validate_json(write_route.calls.last.request.content) == { + "data": { + "password": "new-team-secret-value", + "description": "Rotated from team-old-secret", + } + } + assert new_route.call_count == 1 + assert delete_route.call_count == 1 + + +@pytest.mark.asyncio +@respx.mock +async def test_hashicorp_secret_manager_rotate_secret_value_mismatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + manager: Final = _build_manager(monkeypatch, {}) + respx.post(f"{VAULT_ADDR}/v1/auth/approle/login").respond(json=LOGIN_RESPONSE) + respx.get(f"{VAULT_ADDR}/v1/secret/data/old-secret").respond( + json={"data": {"data": {"key": "old-secret-value"}}} + ) + respx.post(f"{VAULT_ADDR}/v1/secret/data/new-secret").respond( + json={"data": {"version": 1}} + ) + respx.get(f"{VAULT_ADDR}/v1/secret/data/new-secret").respond( + json={"data": {"data": {"key": "different-value"}}} + ) + + response: Final = await manager.async_rotate_secret( + "old-secret", "new-secret", "expected-value" + ) + + assert response["status"] == "error" + assert "mismatch" in response["message"].lower() + assert "expected-value" in response["message"] with Path(__file__).with_name("hashicorp_vault_parity.json").open() as parity_file: