diff --git a/tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py b/tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py new file mode 100644 index 0000000000..b72ff75002 --- /dev/null +++ b/tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py @@ -0,0 +1,129 @@ +import json +import os +import sys + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + + +import pytest +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from fastapi.testclient import TestClient + +import litellm +from litellm.proxy._types import SpecialHeaders +from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware + + +# Fake auth functions to simulate valid and invalid auth behavior. +async def fake_valid_auth(request, api_key): + # Simulate valid authentication: do nothing (i.e. pass) + return + + +async def fake_invalid_auth(request, api_key): + print("running fake invalid auth", request, api_key) + # Simulate invalid auth by raising an exception. + raise Exception("Invalid API key") + + +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + +@pytest.fixture +def app_with_middleware(): + """Create a FastAPI app with the PrometheusAuthMiddleware and dummy endpoints.""" + app = FastAPI() + # Add the PrometheusAuthMiddleware to the app. + app.add_middleware(PrometheusAuthMiddleware) + + @app.get("/metrics") + async def metrics(): + return {"msg": "metrics OK"} + + # Also allow /metrics/ (trailing slash) + @app.get("/metrics/") + async def metrics_slash(): + return {"msg": "metrics OK"} + + @app.get("/chat/completions") + async def chat(): + return {"msg": "chat completions OK"} + + @app.get("/embeddings") + async def embeddings(): + return {"msg": "embeddings OK"} + + return app + + +def test_valid_auth_metrics(app_with_middleware, monkeypatch): + """ + Test that a request to /metrics (and /metrics/) with valid auth headers passes. + """ + # Enable auth on metrics endpoints. + litellm.require_auth_for_metrics_endpoint = True + # Patch the auth function to simulate a valid authentication. + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + fake_valid_auth, + ) + + client = TestClient(app_with_middleware) + headers = {SpecialHeaders.openai_authorization.value: "valid"} + + # Test for /metrics (no trailing slash) + response = client.get("/metrics", headers=headers) + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"} + + # Test for /metrics/ (with trailing slash) + response = client.get("/metrics/", headers=headers) + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"} + + +def test_invalid_auth_metrics(app_with_middleware, monkeypatch): + """ + Test that a request to /metrics with invalid auth headers fails with a 401. + """ + litellm.require_auth_for_metrics_endpoint = True + # Patch the auth function to simulate a failed authentication. + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + fake_invalid_auth, + ) + + client = TestClient(app_with_middleware) + headers = {SpecialHeaders.openai_authorization.value: "invalid"} + + response = client.get("/metrics", headers=headers) + assert response.status_code == 401, response.text + assert "Unauthorized access to metrics endpoint" in response.text + + +def test_no_auth_metrics_when_disabled(app_with_middleware, monkeypatch): + """ + Test that when require_auth_for_metrics_endpoint is False, requests to /metrics + bypass the auth check. + """ + litellm.require_auth_for_metrics_endpoint = False + + # To ensure auth is not run, patch the auth function with one that will raise if called. + def should_not_be_called(*args, **kwargs): + raise Exception("Auth should not be called") + + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + should_not_be_called, + ) + + client = TestClient(app_with_middleware) + response = client.get("/metrics") + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"}