From f2b5a5f8323e37161bee35f885c780973e8e5f41 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 31 Oct 2025 20:50:31 -0700 Subject: [PATCH] test google endpoints --- .../proxy/common_utils/custom_openapi_spec.py | 5 +++ .../test_google_api_endpoints.py | 34 ++++++++++++++++--- 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index 5bd8534a9a6..d57dde04590 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -49,6 +49,11 @@ class CustomOpenAPISpec: except AttributeError: # If both methods fail, return None return None + except Exception as e: + # FastAPI 0.120+ may fail schema generation for certain types (e.g., openai.Timeout) + # Log the error and return None to skip schema generation for this model + verbose_proxy_logger.debug(f"Failed to generate schema for {model_class}: {e}") + return None @staticmethod def add_schema_to_components(openapi_schema: Dict[str, Any], schema_name: str, schema_def: Dict[str, Any]) -> None: diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 6e1c6e11be0..923cb2c6fb8 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -20,13 +20,19 @@ def test_google_generate_content_endpoint(): """Test that the google_generate_content endpoint correctly routes requests""" # Skip this test if we can't import the required modules due to missing dependencies try: + from fastapi import FastAPI from fastapi.testclient import TestClient + from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") + # Create a FastAPI app and include the router (required for FastAPI 0.120+) + app = FastAPI() + app.include_router(google_router) + # Create a test client - client = TestClient(google_router) + client = TestClient(app) # Mock the router's agenerate_content method with patch("litellm.proxy.proxy_server.llm_router") as mock_router: @@ -52,13 +58,19 @@ def test_google_stream_generate_content_endpoint(): """Test that the google_stream_generate_content endpoint correctly routes streaming requests""" # Skip this test if we can't import the required modules due to missing dependencies try: + from fastapi import FastAPI from fastapi.testclient import TestClient + from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") + # Create a FastAPI app and include the router (required for FastAPI 0.120+) + app = FastAPI() + app.include_router(google_router) + # Create a test client - client = TestClient(google_router) + client = TestClient(app) # Mock the router's agenerate_content_stream method to return a stream async def mock_stream_generator(): @@ -91,14 +103,20 @@ def test_google_stream_generate_content_endpoint(): def test_google_generate_content_with_cost_tracking_metadata(): """Test that the google_generate_content endpoint includes user metadata for cost tracking""" try: + from fastapi import FastAPI from fastapi.testclient import TestClient - from litellm.proxy.google_endpoints.endpoints import router as google_router + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") + # Create a FastAPI app and include the router (required for FastAPI 0.120+) + app = FastAPI() + app.include_router(google_router) + # Create a test client - client = TestClient(google_router) + client = TestClient(app) # Mock all required proxy server dependencies with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ @@ -150,13 +168,19 @@ def test_google_generate_content_with_cost_tracking_metadata(): def test_google_stream_generate_content_with_cost_tracking_metadata(): """Test that the google_stream_generate_content endpoint includes user metadata for cost tracking""" try: + from fastapi import FastAPI from fastapi.testclient import TestClient + from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") + # Create a FastAPI app and include the router (required for FastAPI 0.120+) + app = FastAPI() + app.include_router(google_router) + # Create a test client - client = TestClient(google_router) + client = TestClient(app) # Mock the router's agenerate_content_stream method to return a stream mock_stream = AsyncMock()