From b72a03050140e8514ccefdab15ea9b277569ba60 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:56:52 -0700 Subject: [PATCH] test: take keys out of the legacy proxy, enterprise and mcp unit tests before moving them (#42901) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .github/workflows/test-unit-proxy-db.yml | 2 +- .../test_token_counter_gemini_contents_e2e.py | 78 ++++ .../test_prometheus_unit_tests.py | 10 +- .../routing/test_user_config_routing.py | 103 +++++ .../mcp_tests/test_aresponses_api_with_mcp.py | 377 +---------------- .../test_aresponses_api_with_mcp_providers.py | 389 ++++++++++++++++++ .../test_proxy_custom_auth.py | 5 +- .../test_proxy_pass_user_config.py | 114 ----- tests/proxy_unit_tests/test_proxy_server.py | 53 --- .../test_proxy_server_gemini_pass_through.py | 51 +++ .../test_proxy_token_counter.py | 159 +------ tests/proxy_unit_tests/test_proxy_utils.py | 4 +- 12 files changed, 652 insertions(+), 693 deletions(-) create mode 100644 tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py create mode 100644 tests/integration/routing/test_user_config_routing.py create mode 100644 tests/mcp_tests/test_aresponses_api_with_mcp_providers.py delete mode 100644 tests/proxy_unit_tests/test_proxy_pass_user_config.py create mode 100644 tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 4af7a161984..73015ac6e02 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -106,6 +106,7 @@ jobs: - test-group: proxy-server-core test-path: >- tests/proxy_unit_tests/test_proxy_server.py + tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py tests/proxy_unit_tests/test_aproxy_startup.py workers: 4 dist: loadscope @@ -115,7 +116,6 @@ jobs: tests/proxy_unit_tests/test_proxy_config_unit_test.py tests/proxy_unit_tests/test_proxy_routes.py tests/proxy_unit_tests/test_server_root_path.py - tests/proxy_unit_tests/test_proxy_pass_user_config.py tests/proxy_unit_tests/test_proxy_token_counter.py tests/proxy_unit_tests/test_request_size_limit_middleware.py tests/proxy_unit_tests/test_multipart_bypass_repro.py diff --git a/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py b/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py new file mode 100644 index 00000000000..b2b8f46ab0a --- /dev/null +++ b/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py @@ -0,0 +1,78 @@ +"""Live e2e: `/utils/token_counter?call_endpoint=true` counts Gemini `contents` upstream. + +Google's countTokens API is the only tokenizer that knows Gemini's real token +boundaries, so the proxy must forward `contents` to it for both the AI Studio and +Vertex deployments and hand back the provider's `promptTokensDetails`. Claude on +Vertex is covered by `/v1/messages/count_tokens`; this is the Gemini `contents` +route the claude_code rows never reach +""" + +from __future__ import annotations + +import pytest +from e2e_config import unique_marker +from e2e_http import require_successful_call +from proxy_client import ProxyClient +from pydantic import BaseModel + +pytestmark = pytest.mark.e2e + +GEMINI_DEPLOYMENTS = ("gemini-2.5-flash", "gemini-2.5-flash-vertex") + + +class _Part(BaseModel): + text: str + + +class _Content(BaseModel): + parts: tuple[_Part, ...] + + +class _TokenCountBody(BaseModel): + model: str + contents: tuple[_Content, ...] + + +class _CallEndpoint(BaseModel): + call_endpoint: bool = True + + +class _ModalityTokens(BaseModel): + modality: str + tokenCount: int + + +class _CountTokensUpstream(BaseModel): + totalTokens: int + promptTokensDetails: tuple[_ModalityTokens, ...] + + +class _TokenCountResponse(BaseModel): + total_tokens: int + request_model: str + model_used: str + tokenizer_type: str + original_response: _CountTokensUpstream + + +class TestGeminiContentsTokenCounting: + @pytest.mark.parametrize("model", GEMINI_DEPLOYMENTS) + def test_contents_are_counted_by_the_provider_endpoint( + self, proxy: ProxyClient, scoped_key: str, model: str + ) -> None: + text = f"Hello world, how are you doing today? {unique_marker()}" + body = _TokenCountBody(model=model, contents=(_Content(parts=(_Part(text=text),)),)) + + result = proxy.transport.send( + "/utils/token_counter", + headers=proxy.transport.bearer(scoped_key), + json=body, + params=_CallEndpoint(), + ) + + require_successful_call(result) + counted = _TokenCountResponse.model_validate_json(result.body) + assert counted.request_model == model, counted + assert counted.original_response.totalTokens == counted.total_tokens > 0, counted + assert counted.original_response.promptTokensDetails, counted + assert all(detail.tokenCount > 0 for detail in counted.original_response.promptTokensDetails), counted diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py index 28fd03daf37..5b26dad269e 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py @@ -13,7 +13,6 @@ import asyncio from dotenv import load_dotenv load_dotenv() -import os from unittest.mock import MagicMock @@ -165,9 +164,9 @@ async def test_prometheus_metric_tracking(): "model_name": "gpt-5-mini", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_AI_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": "sk-azure-unit-test", + "api_version": "2025-01-01-preview", + "api_base": "https://unit-test.openai.azure.com", }, "model_info": {"id": "azure-model-id"}, }, @@ -180,9 +179,6 @@ async def test_prometheus_metric_tracking(): }, ], provider_budget_config=provider_budget_config, - redis_host=os.getenv("REDIS_HOST"), - redis_port=int(os.getenv("REDIS_PORT", 6379)), - redis_password=os.getenv("REDIS_PASSWORD"), ) try: diff --git a/tests/integration/routing/test_user_config_routing.py b/tests/integration/routing/test_user_config_routing.py new file mode 100644 index 00000000000..1c50a78201b --- /dev/null +++ b/tests/integration/routing/test_user_config_routing.py @@ -0,0 +1,103 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +USER_KEY: Final = "sk-user-supplied-" + uuid.uuid4().hex + + +def _completion(request: Request) -> Reply: + if request.target != "/v1/chat/completions": + return Reply(status=404, body=b"{}") + body: Final = json.loads(request.body) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-user-config", + "object": "chat.completion", + "created": 0, + "model": body["model"], + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _user_config(upstream_url: str) -> dict[str, object]: + return { + "model_list": [ + { + "model_name": "user-config-deployment", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_base": upstream_url + "/v1", + "api_key": USER_KEY, + }, + } + ], + "num_retries": 0, + } + + +def _opt_in_config(directory: Path, upstream_url: str) -> Path: + config: Final = directory / "allow_client_side_credentials_config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "admin-deployment", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_base": upstream_url + "/v1", + "api_key": "sk-admin-configured", + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + "allow_client_side_credentials": True, + }, + } + ) + ) + return config + + +def _request_body(upstream_url: str) -> dict[str, object]: + return { + "model": "user-config-deployment", + "messages": [{"role": "user", "content": "user config control"}], + "user_config": _user_config(upstream_url), + } + + +def test_user_config_routes_to_the_user_supplied_deployment_when_opted_in(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(_completion) as upstream: + config: Final = _opt_in_config(tmp_path, upstream.url) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + response: Final = candidate.request("POST", "/v1/chat/completions", _request_body(upstream.url)) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "routed" + outbound: Final = tuple(upstream.received.get_nowait() for _ in range(upstream.received.qsize())) + completions: Final = tuple(request for request in outbound if request.target == "/v1/chat/completions") + assert len(completions) == 1, outbound + assert completions[0].headers["authorization"] == f"Bearer {USER_KEY}" + assert json.loads(completions[0].body)["model"] == "gpt-4.1-mini" + + +def test_user_config_is_rejected_without_the_opt_in(gateway: Gateway) -> None: + with wire_server(_completion) as upstream: + response: Final = gateway.request("POST", "/v1/chat/completions", _request_body(upstream.url)) + assert response.status_code == 401, response.text + assert "user_config is not allowed in request body" in response.text + assert upstream.received.empty() diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index eb6f78b57a1..de0dc78af43 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -1,5 +1,3 @@ -import logging -import os import pytest from mcp.types import Tool as MCPTool from typing import List, Any, cast @@ -846,161 +844,6 @@ async def test_streaming_mcp_events_validation(): assert mock_get_tools.called, "MCP tools should have been fetched" -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", - [ - pytest.param("gpt-4o-mini", id="openai"), - pytest.param("claude-haiku-4-5", id="anthropic"), - ], -) -async def test_streaming_responses_api_with_mcp_tools( - model: str, caplog: pytest.LogCaptureFixture -): - """ - Test the streaming responses API with MCP tools when using server_url="litellm_proxy" - - Under the hood the follow occurs - - - MCP: responses called litellm MCP manager.list_tools (MOCKED) - - Request 1: Made to model under test with fetched tools (REAL LLM CALL) - - MCP: Execute tool call from request 1 and returns result (MOCKED) - - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) - - Return the user the result of request 2 - """ - # Skip test if API keys are not set for the respective models - if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv( - "ANTHROPIC_API_KEY" - ): - pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test") - if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( - "OPENAI_API_KEY" - ): - pytest.skip("OPENAI_API_KEY not set, skipping openai model test") - - from unittest.mock import AsyncMock, patch - - print("๐Ÿงช Testing basic streaming with MCP tools...") - - # Mock MCP tools that would be returned from the manager - mock_mcp_tools = [ - MCPTool.model_validate({ - "name": "search_repo", - "description": "Search BerriAI/litellm repository for information", - "inputSchema": { - "type": "object", - "properties": { - "query": {"type": "string", "description": "Search query"} - }, - "required": ["query"], - }, - }, by_name=False) - ] - - # Only mock the MCP-specific operations, let LLM responses be real - with caplog.at_level(logging.ERROR): - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_get_mcp_tools_from_manager", - new_callable=AsyncMock, - ) as mock_get_tools, - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_execute_tool_calls", - new_callable=AsyncMock, - ) as mock_execute_tools, - ): - # Setup MCP mocks only - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - # Create a dynamic mock that will match the actual tool call ID from the LLM response - def mock_execute_tool_calls_side_effect( - tool_calls, user_api_key_auth, **kwargs - ): - """Mock function that returns results matching the actual tool call IDs from the LLM""" - results = [] - for tool_call in tool_calls: - # Extract call_id from the tool call - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, "call_id"): - call_id = tool_call.call_id - elif hasattr(tool_call, "id"): - call_id = tool_call.id - - if call_id: - results.append( - { - "tool_call_id": call_id, - "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", - } - ) - return results - - mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect - - # Make the actual call - LLM responses will be real - mcp_tool_config = cast( - Any, - { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never", - }, - ) - response = await litellm.aresponses( - model=model, - tools=[mcp_tool_config], - tool_choice="required", - input=[ - { - "role": "user", - "type": "message", - "content": "give me a TLDR of what BerriAI/litellm is about", - } - ], - stream=True, - ) - - print(f"๐Ÿ“‹ Response type: {type(response)}") - assert hasattr( - response, "__aiter__" - ), "Response should be an async streaming response" - - # Collect streaming chunks - chunks = [] - async for chunk in response: - chunks.append(chunk) - print(f"๐Ÿ“ฆ Chunk type: {getattr(chunk, 'type', 'unknown')}") - - print(f"๐Ÿ“Š Total chunks received: {len(chunks)}") - - # Verify MCP mocks were called (may be called multiple times in streaming) - assert ( - mock_get_tools.call_count >= 1 - ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" - print(f"MCP tools fetched: {len(mock_mcp_tools)}") - - # Verify we got a response - assert response is not None - assert len(chunks) > 0, "Should have received streaming chunks" - - print("Basic streaming responses API with MCP tools test passed!") - - lite_errors = [ - record - for record in caplog.records - if record.levelno >= logging.ERROR - and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) - ] - assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( - record.getMessage() for record in lite_errors - ) - - @pytest.mark.asyncio async def test_mcp_parameter_preparation_helpers(): """ @@ -1215,7 +1058,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): The test mocks the MCP manager response but validates the actual tools sent to the LLM to ensure no duplication occurs. """ - from unittest.mock import AsyncMock, patch, call + from unittest.mock import AsyncMock, patch from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -1432,221 +1275,3 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): "tools_per_call": [len(tools) for tools in llm_call_tools], "duplicate_tools_found": False, } - - -@pytest.mark.asyncio -@pytest.mark.parametrize("model", ["gpt-4o-mini"]) -async def test_streaming_mcp_event_order_and_response_id_consistency( - model: str, caplog: pytest.LogCaptureFixture -): - """ - Test that: - 1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events) - 2. All response lifecycle events share the same response ID within a cycle - """ - if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( - "OPENAI_API_KEY" - ): - pytest.skip("OPENAI_API_KEY not set, skipping openai model test") - - from unittest.mock import AsyncMock, patch - - mock_mcp_tools = [ - MCPTool.model_validate({ - "name": "get_weather", - "description": "Get weather for a city", - "inputSchema": { - "type": "object", - "properties": { - "city": {"type": "string", "description": "City name"} - }, - "required": ["city"], - }, - }, by_name=False) - ] - - with caplog.at_level(logging.ERROR): - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_get_mcp_tools_from_manager", - new_callable=AsyncMock, - ) as mock_get_tools, - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_execute_tool_calls", - new_callable=AsyncMock, - ) as mock_execute_tools, - ): - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs): - results = [] - for tool_call in tool_calls: - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, "call_id"): - call_id = tool_call.call_id - elif hasattr(tool_call, "id"): - call_id = tool_call.id - if call_id: - results.append( - { - "tool_call_id": call_id, - "result": "Sunny, 72ยฐF", - } - ) - return results - - mock_execute_tools.side_effect = mock_execute_side_effect - - mcp_tool_config = cast( - Any, - { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never", - }, - ) - - response = await litellm.aresponses( - model=model, - tools=[mcp_tool_config], - input=[ - { - "role": "user", - "type": "message", - "content": "What's the weather in San Francisco?", - } - ], - stream=True, - ) - - events = [] - async for chunk in response: - events.append(chunk) - - assert len(events) > 0, "Should receive streaming events" - - created_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.created" - ), - None, - ) - in_progress_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.in_progress" - ), - None, - ) - output_item_added_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.output_item.added" - ), - None, - ) - mcp_in_progress_idx = next( - ( - i - for i, e in enumerate(events) - if "mcp_list_tools.in_progress" in str(getattr(e, "type", "")) - ), - None, - ) - completed_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.completed" - ), - None, - ) - - assert created_idx is not None, "response.created event should be present" - assert ( - in_progress_idx is not None - ), "response.in_progress event should be present" - assert ( - output_item_added_idx is not None - ), "response.output_item.added event should be present" - - assert ( - created_idx < in_progress_idx - ), "response.created should come before response.in_progress" - assert ( - in_progress_idx < output_item_added_idx - ), "response.in_progress should come before response.output_item.added" - - if mcp_in_progress_idx is not None: - assert ( - output_item_added_idx < mcp_in_progress_idx - ), "response.output_item.added should come before response.mcp_list_tools.in_progress" - - response_ids = [] - for i, event in enumerate(events): - event_type = getattr(event, "type", None) - if hasattr(event, "response"): - response_obj = getattr(event, "response", None) - if response_obj and hasattr(response_obj, "id"): - event_type_value = ( - event_type.value - if hasattr(event_type, "value") - else str(event_type) - ) - if any( - x in event_type_value - for x in [ - "response.created", - "response.in_progress", - "response.completed", - ] - ): - response_ids.append((i, event_type_value, response_obj.id)) - - assert ( - len(response_ids) >= 2 - ), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}" - - cycles = [] - current_cycle = [] - current_id = None - - for idx, event_type, resp_id in response_ids: - if current_id is None or resp_id == current_id: - current_cycle.append((idx, event_type, resp_id)) - current_id = resp_id - else: - if current_cycle: - cycles.append(current_cycle) - current_cycle = [(idx, event_type, resp_id)] - current_id = resp_id - if current_cycle: - cycles.append(current_cycle) - - for cycle_num, cycle in enumerate(cycles): - cycle_ids = set(resp_id for _, _, resp_id in cycle) - assert ( - len(cycle_ids) == 1 - ), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs" - - assert ( - completed_idx is not None - ), "response.completed event should be present" - - lite_errors = [ - record - for record in caplog.records - if record.levelno >= logging.ERROR - and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) - ] - assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( - record.getMessage() for record in lite_errors - ) diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py b/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py new file mode 100644 index 00000000000..72a0415cf88 --- /dev/null +++ b/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py @@ -0,0 +1,389 @@ +import logging +import os +import pytest +from mcp.types import Tool as MCPTool +from typing import Any, cast + +import litellm +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + + +class MockUserAPIKeyAuth: + """Mock UserAPIKeyAuth for testing""" + + def __init__(self): + self.api_key = "test_key" + self.user_id = "test_user" + self.team_id = "test_team" + self.user_email = "test@example.com" + self.max_budget = 100.0 + self.spend = 0.0 + self.models = [] + self.aliases = {} + self.config = {} + self.permissions = {} + self.metadata = {} + self.object_permission_id = "test_permission_id" + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + pytest.param("gpt-4o-mini", id="openai"), + pytest.param("claude-haiku-4-5", id="anthropic"), + ], +) +async def test_streaming_responses_api_with_mcp_tools( + model: str, caplog: pytest.LogCaptureFixture +): + """ + Test the streaming responses API with MCP tools when using server_url="litellm_proxy" + + Under the hood the follow occurs + + - MCP: responses called litellm MCP manager.list_tools (MOCKED) + - Request 1: Made to model under test with fetched tools (REAL LLM CALL) + - MCP: Execute tool call from request 1 and returns result (MOCKED) + - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) + + Return the user the result of request 2 + """ + if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv( + "ANTHROPIC_API_KEY" + ): + pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test") + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( + "OPENAI_API_KEY" + ): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") + + from unittest.mock import AsyncMock, patch + + print("๐Ÿงช Testing basic streaming with MCP tools...") + + mock_mcp_tools = [ + MCPTool.model_validate({ + "name": "search_repo", + "description": "Search BerriAI/litellm repository for information", + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"} + }, + "required": ["query"], + }, + }, by_name=False) + ] + + with caplog.at_level(logging.ERROR): + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_get_mcp_tools_from_manager", + new_callable=AsyncMock, + ) as mock_get_tools, + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + new_callable=AsyncMock, + ) as mock_execute_tools, + ): + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + def mock_execute_tool_calls_side_effect( + tool_calls, user_api_key_auth, **kwargs + ): + """Mock function that returns results matching the actual tool call IDs from the LLM""" + results = [] + for tool_call in tool_calls: + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, "call_id"): + call_id = tool_call.call_id + elif hasattr(tool_call, "id"): + call_id = tool_call.id + + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect + + mcp_tool_config = cast( + Any, + { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + tool_choice="required", + input=[ + { + "role": "user", + "type": "message", + "content": "give me a TLDR of what BerriAI/litellm is about", + } + ], + stream=True, + ) + + print(f"๐Ÿ“‹ Response type: {type(response)}") + assert hasattr( + response, "__aiter__" + ), "Response should be an async streaming response" + + chunks = [] + async for chunk in response: + chunks.append(chunk) + print(f"๐Ÿ“ฆ Chunk type: {getattr(chunk, 'type', 'unknown')}") + + print(f"๐Ÿ“Š Total chunks received: {len(chunks)}") + + assert ( + mock_get_tools.call_count >= 1 + ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" + print(f"MCP tools fetched: {len(mock_mcp_tools)}") + + assert response is not None + assert len(chunks) > 0, "Should have received streaming chunks" + + print("Basic streaming responses API with MCP tools test passed!") + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) + + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-4o-mini"]) +async def test_streaming_mcp_event_order_and_response_id_consistency( + model: str, caplog: pytest.LogCaptureFixture +): + """ + Test that: + 1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events) + 2. All response lifecycle events share the same response ID within a cycle + """ + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( + "OPENAI_API_KEY" + ): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") + + from unittest.mock import AsyncMock, patch + + mock_mcp_tools = [ + MCPTool.model_validate({ + "name": "get_weather", + "description": "Get weather for a city", + "inputSchema": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"} + }, + "required": ["city"], + }, + }, by_name=False) + ] + + with caplog.at_level(logging.ERROR): + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_get_mcp_tools_from_manager", + new_callable=AsyncMock, + ) as mock_get_tools, + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + new_callable=AsyncMock, + ) as mock_execute_tools, + ): + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs): + results = [] + for tool_call in tool_calls: + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, "call_id"): + call_id = tool_call.call_id + elif hasattr(tool_call, "id"): + call_id = tool_call.id + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "Sunny, 72ยฐF", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_side_effect + + mcp_tool_config = cast( + Any, + { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + input=[ + { + "role": "user", + "type": "message", + "content": "What's the weather in San Francisco?", + } + ], + stream=True, + ) + + events = [] + async for chunk in response: + events.append(chunk) + + assert len(events) > 0, "Should receive streaming events" + + created_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.created" + ), + None, + ) + in_progress_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.in_progress" + ), + None, + ) + output_item_added_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.output_item.added" + ), + None, + ) + mcp_in_progress_idx = next( + ( + i + for i, e in enumerate(events) + if "mcp_list_tools.in_progress" in str(getattr(e, "type", "")) + ), + None, + ) + completed_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.completed" + ), + None, + ) + + assert created_idx is not None, "response.created event should be present" + assert ( + in_progress_idx is not None + ), "response.in_progress event should be present" + assert ( + output_item_added_idx is not None + ), "response.output_item.added event should be present" + + assert ( + created_idx < in_progress_idx + ), "response.created should come before response.in_progress" + assert ( + in_progress_idx < output_item_added_idx + ), "response.in_progress should come before response.output_item.added" + + if mcp_in_progress_idx is not None: + assert ( + output_item_added_idx < mcp_in_progress_idx + ), "response.output_item.added should come before response.mcp_list_tools.in_progress" + + response_ids = [] + for i, event in enumerate(events): + event_type = getattr(event, "type", None) + if hasattr(event, "response"): + response_obj = getattr(event, "response", None) + if response_obj and hasattr(response_obj, "id"): + event_type_value = ( + event_type.value + if hasattr(event_type, "value") + else str(event_type) + ) + if any( + x in event_type_value + for x in [ + "response.created", + "response.in_progress", + "response.completed", + ] + ): + response_ids.append((i, event_type_value, response_obj.id)) + + assert ( + len(response_ids) >= 2 + ), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}" + + cycles = [] + current_cycle = [] + current_id = None + + for idx, event_type, resp_id in response_ids: + if current_id is None or resp_id == current_id: + current_cycle.append((idx, event_type, resp_id)) + current_id = resp_id + else: + if current_cycle: + cycles.append(current_cycle) + current_cycle = [(idx, event_type, resp_id)] + current_id = resp_id + if current_cycle: + cycles.append(current_cycle) + + for cycle_num, cycle in enumerate(cycles): + cycle_ids = set(resp_id for _, _, resp_id in cycle) + assert ( + len(cycle_ids) == 1 + ), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs" + + assert ( + completed_idx is not None + ), "response.completed event should be present" + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) diff --git a/tests/proxy_unit_tests/test_proxy_custom_auth.py b/tests/proxy_unit_tests/test_proxy_custom_auth.py index b575e4c85c6..dbbad0dab1e 100644 --- a/tests/proxy_unit_tests/test_proxy_custom_auth.py +++ b/tests/proxy_unit_tests/test_proxy_custom_auth.py @@ -53,8 +53,7 @@ def test_custom_auth(client): "max_tokens": 10, } # Your bearer token - token = os.getenv("PROXY_MASTER_KEY") - print(f"token: {token}") + token = "sk-unit-test-master" headers = {"Authorization": f"Bearer {token}"} with pytest.raises(Exception, match="Authentication Error, Failed custom auth") as exc_info: client.post("/chat/completions", json=test_data, headers=headers) @@ -71,7 +70,7 @@ def test_custom_auth_bearer(client): "max_tokens": 10, } # Your bearer token - token = os.getenv("PROXY_MASTER_KEY") + token = "sk-unit-test-master" headers = {"Authorization": f"WITHOUT BEAR Er {token}"} with pytest.raises(Exception, match="CustomAuth - Malformed API Key passed in") as exc_info: diff --git a/tests/proxy_unit_tests/test_proxy_pass_user_config.py b/tests/proxy_unit_tests/test_proxy_pass_user_config.py deleted file mode 100644 index 91911c142ea..00000000000 --- a/tests/proxy_unit_tests/test_proxy_pass_user_config.py +++ /dev/null @@ -1,114 +0,0 @@ -import sys, os -import traceback -from dotenv import load_dotenv - -load_dotenv() -import io - -# this file is to test litellm/proxy - -import pytest, logging, asyncio -import litellm -from litellm import embedding, completion, completion_cost, Timeout -from litellm import RateLimitError - -# Configure logging -logging.basicConfig( - level=logging.DEBUG, # Set the desired logging level - format="%(asctime)s - %(levelname)s - %(message)s", -) - -# test /chat/completion request to the proxy -from fastapi.testclient import TestClient -from fastapi import FastAPI -from litellm.proxy.proxy_server import ( - router, - save_worker_config, - initialize, -) # Replace with the actual module where your FastAPI router is defined - -# Your bearer token -token = "sk-1234" - -headers = {"Authorization": f"Bearer {token}"} - - -@pytest.fixture(scope="function") -def client_no_auth(): - # Assuming litellm.proxy.proxy_server is an object - from litellm.proxy.proxy_server import cleanup_router_config_variables - - cleanup_router_config_variables() - filepath = os.path.dirname(os.path.abspath(__file__)) - config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" - # initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables - asyncio.run(initialize(config=config_fp, debug=True)) - app = FastAPI() - app.include_router(router) # Include your router in the test app - - return TestClient(app) - - -@pytest.mark.skipif( - os.environ.get("AZURE_AI_API_KEY") is None - or os.environ.get("OPENAI_API_KEY") is None, - reason="AZURE_AI_API_KEY or OPENAI_API_KEY not set - skipping integration test", -) -def test_chat_completion(client_no_auth): - global headers - - from litellm.types.router import RouterConfig, ModelConfig - from litellm.types.completion import CompletionRequest - - user_config = RouterConfig( - model_list=[ - ModelConfig( - model_name="user-azure-instance", - litellm_params=CompletionRequest( - model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_AI_API_KEY"), - api_version=os.getenv("AZURE_API_VERSION"), - api_base=os.getenv("AZURE_AI_API_BASE"), - timeout=10, - ), - tpm=240000, - rpm=1800, - ), - ModelConfig( - model_name="user-openai-instance", - litellm_params=CompletionRequest( - model="gpt-3.5-turbo", - api_key=os.getenv("OPENAI_API_KEY"), - timeout=10, - ), - tpm=240000, - rpm=1800, - ), - ], - num_retries=2, - allowed_fails=3, - fallbacks=[{"user-azure-instance": ["user-openai-instance"]}], - ).dict() - - try: - # Your test data - test_data = { - "model": "user-azure-instance", - "messages": [ - {"role": "user", "content": "hi"}, - ], - "max_tokens": 10, - "user_config": user_config, - } - - print("testing proxy server with chat completions") - response = client_no_auth.post("/v1/chat/completions", json=test_data) - print(f"response - {response.text}") - assert response.status_code == 200 - result = response.json() - print(f"Received response: {result}") - except Exception as e: - pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") - - -# Run the test diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index ed0380058a5..5be27b3ad72 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2065,59 +2065,6 @@ async def test_add_callback_via_key_litellm_pre_call_utils_langsmith( assert new_data["failure_callback"] == expected_failure_callbacks -@pytest.mark.skipif( - not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"), - reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.", -) -@pytest.mark.asyncio -async def test_gemini_pass_through_endpoint(): - from starlette.datastructures import URL - - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - Request, - Response, - gemini_proxy_route, - ) - - body = b""" - { - "contents": [{ - "parts":[{ - "text": "The quick brown fox jumps over the lazy dog." - }] - }] - } - """ - - # Construct the scope dictionary - scope = { - "type": "http", - "method": "POST", - "path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens", - "query_string": b"key=sk-1234", - "headers": [ - (b"content-type", b"application/json"), - ], - } - - # Create a new Request object - async def async_receive(): - return {"type": "http.request", "body": body, "more_body": False} - - request = Request( - scope=scope, - receive=async_receive, - ) - - resp = await gemini_proxy_route( - endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234", - request=request, - fastapi_response=Response(), - ) - - print(resp.body) - - @pytest.mark.parametrize("hidden", [True, False]) @pytest.mark.asyncio async def test_model_info_alias_without_prisma(hidden): diff --git a/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py b/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py new file mode 100644 index 00000000000..2453ec3bfe3 --- /dev/null +++ b/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py @@ -0,0 +1,51 @@ +import os + +import pytest + + +@pytest.mark.skipif( + not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"), + reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.", +) +@pytest.mark.asyncio +async def test_gemini_pass_through_endpoint(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + Request, + Response, + gemini_proxy_route, + ) + + body = b""" + { + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }] + }] + } + """ + + scope = { + "type": "http", + "method": "POST", + "path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens", + "query_string": b"key=sk-1234", + "headers": [ + (b"content-type", b"application/json"), + ], + } + + async def async_receive(): + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + scope=scope, + receive=async_receive, + ) + + await gemini_proxy_route( + endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234", + request=request, + fastapi_response=Response(), + ) + diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 39ec4bb1887..8590e959961 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -2,10 +2,7 @@ # 1. Generate a Key, and use it to make a call -import json import logging -import os -import tempfile from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -35,79 +32,6 @@ from litellm.types.utils import TokenCountResponse verbose_proxy_logger.setLevel(level=logging.DEBUG) -def get_vertex_ai_creds_json() -> dict: - # 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 - - return service_account_key_data - - -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) - - @pytest.mark.asyncio async def test_vLLM_token_counting(): """ @@ -223,10 +147,12 @@ async def test_anthropic_messages_count_tokens_endpoint(): - Should return response in Anthropic format: {"input_tokens": } - Should work as wrapper around internal token_counter function """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request from unittest.mock import MagicMock + from fastapi import Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_request_data = { @@ -295,10 +221,12 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): - Should still work and return Anthropic format - Should call internal token_counter with from_anthropic_endpoint=True """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request from unittest.mock import MagicMock + from fastapi import Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_request_data = { @@ -435,10 +363,12 @@ async def test_anthropic_endpoint_error_handling(): """ Test error handling in the /v1/messages/count_tokens endpoint """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request, HTTPException from unittest.mock import MagicMock + from fastapi import HTTPException, Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_user_api_key_dict = MagicMock() @@ -474,8 +404,10 @@ async def test_anthropic_endpoint_error_handling(): @pytest.mark.asyncio async def test_factory_anthropic_endpoint_calls_anthropic_counter(): """Test that /v1/messages/count_tokens with Anthropic model uses Anthropic counter.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -531,8 +463,10 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): @pytest.mark.asyncio async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): """Test that /v1/messages/count_tokens with GPT-4 does NOT use Anthropic counter.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -590,8 +524,10 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): @pytest.mark.asyncio async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): """Test that /utils/token_counter does NOT use Anthropic counter even with Anthropic model.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -678,57 +614,6 @@ async def test_factory_registration(): assert not counter.should_use_token_counting_api(custom_llm_provider=None) -@pytest.mark.skip( - reason="Requires Google/Vertex AI credentials (GEMINI_API_KEY or VERTEX_AI_PRIVATE_KEY)." -) -@pytest.mark.asyncio -@pytest.mark.parametrize("model_name", ["gemini-2.5-pro", "vertex-ai-gemini-2.5-pro"]) -async def test_vertex_ai_gemini_token_counting_with_contents(model_name): - """ - Test token counting for Vertex AI Gemini model using contents format with call_endpoint=True - """ - load_vertex_ai_credentials() - llm_router = Router( - model_list=[ - { - "model_name": "gemini-2.5-pro", - "litellm_params": { - "model": "gemini/gemini-2.5-pro", - }, - }, - { - "model_name": "vertex-ai-gemini-2.5-pro", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-pro", - }, - }, - ] - ) - - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - - # Test with contents format and call_endpoint=True - response = await token_counter( - request=TokenCountRequest( - model=model_name, - contents=[ - {"parts": [{"text": "Hello world, how are you doing today? i am ij"}]} - ], - ), - call_endpoint=True, - ) - - print("Vertex AI Gemini token counting response:", response) - - # validate we have original response - assert response.original_response is not None - assert response.original_response.get("totalTokens") is not None - assert response.original_response.get("promptTokensDetails") is not None - - prompt_tokens_details = response.original_response.get("promptTokensDetails") - assert prompt_tokens_details is not None - - @pytest.mark.asyncio async def test_bedrock_count_tokens_endpoint(): """ @@ -779,7 +664,7 @@ async def test_vertex_ai_anthropic_token_counting(): This tests the token counting implementation for Vertex AI partner models without making actual API calls. Mocks at the handler level to test the full flow. """ - from unittest.mock import AsyncMock, patch, MagicMock + from unittest.mock import patch # Mock the Vertex AI partner models token counter response mock_token_response = { diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 1134f41a940..ab787bbbe27 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -438,7 +438,7 @@ def test_is_request_body_safe_global_enabled( "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), + "api_key": "sk-openai-unit-test", }, } ] @@ -475,7 +475,7 @@ def test_is_request_body_safe_model_enabled( "model_name": "fireworks_ai/*", "litellm_params": { "model": "fireworks_ai/*", - "api_key": os.getenv("FIREWORKS_API_KEY"), + "api_key": "sk-fireworks-unit-test", "configurable_clientside_auth_params": ( ["api_base"] if allow_client_side_credentials else [] ),