litellm/tests/proxy_unit_tests/test_google_endpoint_routing.py

import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
import yaml

sys.path.insert(0, os.path.abspath("../.."))

from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.google_endpoints.endpoints import google_generate_content
from fastapi import Request, Response
from fastapi.datastructures import Headers
from litellm.proxy.proxy_server import initialize
from litellm.utils import ModelResponse

@pytest.fixture
def mock_user_api_key_dict():
    """Mock user API key dictionary."""
    return UserAPIKeyAuth(
        api_key="test_api_key",
        user_id="test_user_id",
        user_email="test@example.com",
        team_id="test_team_id",
        max_budget=100.0,
        spend=0.0,
        user_role="internal_user",
        allowed_cache_controls=[],
        metadata={},
        tpm_limit=None,
        rpm_limit=None,
    )


@pytest.fixture
def mock_request(request):
    """Create a mock FastAPI request with the sample payload."""
    mock_req = MagicMock(spec=Request)
    mock_req.headers = Headers({"content-type": "application/json"})
    mock_req.method = "POST"
    mock_req.url.path = request.param.get("path")
    
    async def mock_body():
        return json.dumps(request.param.get("payload", {})).encode('utf-8')
    
    mock_req.body = mock_body
    return mock_req


@pytest.fixture  
def mock_response():
    """Create a mock FastAPI response."""
    return MagicMock(spec=Response)


@pytest.mark.asyncio
@pytest.mark.parametrize("mock_request", [{"path": "/v1beta/models/bedrock/claude-sonnet-3.7:generateContent", "payload": {"contents": [{"parts":[{"text": "The quick brown fox jumps over the lazy dog."}]}]}}], indirect=True)
async def test_google_generate_content_with_slashes_in_model_name(
    mock_request, mock_response, mock_user_api_key_dict
):
    """
    Test that the google_generate_content endpoint correctly handles model names with slashes.
    """
    config = {
        "model_list": [
            {
                "model_name": "bedrock/claude-sonnet-3.7",
                "litellm_params": {
                    "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
                },
            }
        ]
    }
    filepath = os.path.dirname(os.path.abspath(__file__))
    config_fp = f"{filepath}/test_config.yaml"
    with open(config_fp, "w") as f:
        yaml.dump(config, f)

    try:
        await initialize(config=config_fp)

        with patch("litellm.proxy.proxy_server.llm_router.agenerate_content", new_callable=AsyncMock) as mock_agenerate_content:
            mock_agenerate_content.return_value = ModelResponse()
            
            await google_generate_content(
                request=mock_request,
                model_name="bedrock/claude-sonnet-3.7",
                fastapi_response=mock_response,
                user_api_key_dict=mock_user_api_key_dict,
            )

            mock_agenerate_content.assert_called_once()
            _, call_kwargs = mock_agenerate_content.call_args
            assert call_kwargs["model"] == "bedrock/claude-sonnet-3.7"
    finally:
        if os.path.exists(config_fp):
            os.remove(config_fp)