import asyncio import importlib import json import os import socket import subprocess import sys from datetime import datetime, timezone from pathlib import Path from unittest import mock from unittest.mock import AsyncMock, MagicMock, mock_open, patch import click import httpx import pytest import yaml from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system-path import litellm from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map example_embedding_result = { "object": "list", "data": [ { "object": "embedding", "index": 0, "embedding": [ -0.006929283495992422, -0.005336422007530928, -4.547132266452536e-05, -0.024047505110502243, -0.006929283495992422, -0.005336422007530928, -4.547132266452536e-05, -0.024047505110502243, -0.006929283495992422, -0.005336422007530928, -4.547132266452536e-05, -0.024047505110502243, ], } ], "model": "text-embedding-3-small", "usage": {"prompt_tokens": 5, "total_tokens": 5}, } def mock_patch_aembedding(): return mock.patch( "litellm.proxy.proxy_server.llm_router.aembedding", return_value=example_embedding_result, ) @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)) return TestClient(app) def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): mock_login_result = {"user_id": "test-user"} mock_prisma_client = MagicMock() mock_authenticate_user = AsyncMock(return_value=mock_login_result) mock_create_ui_token_object = MagicMock(return_value={"user_id": "test-user"}) mock_jwt_encode = MagicMock(return_value="signed-token") monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", mock_authenticate_user, ) monkeypatch.setattr( "litellm.proxy.auth.login_utils.create_ui_token_object", mock_create_ui_token_object, ) monkeypatch.setattr("jwt.encode", mock_jwt_encode) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) client = TestClient(app) response = client.post( "/v2/login", json={"username": "alice", "password": "secret"}, ) assert response.status_code == 200 assert response.json() == { "redirect_url": "http://testserver/ui/?login=success", "token": "signed-token", } assert response.cookies.get("token") == "signed-token" mock_authenticate_user.assert_awaited_once_with( username="alice", password="secret", master_key="test-master-key", prisma_client=mock_prisma_client, ) mock_create_ui_token_object.assert_called_once_with( login_result=mock_login_result, general_settings={}, premium_user=False, ) mock_jwt_encode.assert_called_once_with( {"user_id": "test-user"}, "test-master-key", algorithm="HS256", ) def test_login_v2_returns_json_on_proxy_exception(monkeypatch): """Test that /v2/login returns JSON error when ProxyException is raised""" from litellm.proxy._types import ProxyErrorTypes, ProxyException mock_prisma_client = MagicMock() mock_authenticate_user = AsyncMock( side_effect=ProxyException( message="Invalid credentials", type=ProxyErrorTypes.auth_error, param="password", code=401, ) ) monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", mock_authenticate_user, ) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) client = TestClient(app) response = client.post( "/v2/login", json={"username": "alice", "password": "wrong"}, ) assert response.status_code == 401 assert response.headers["content-type"] == "application/json" data = response.json() assert "error" in data assert data["error"]["message"] == "Invalid credentials" assert data["error"]["type"] == "auth_error" def test_login_v2_returns_json_on_http_exception(monkeypatch): """Test that /v2/login converts HTTPException to JSON error response""" from fastapi import HTTPException mock_prisma_client = MagicMock() mock_authenticate_user = AsyncMock( side_effect=HTTPException(status_code=401, detail="Unauthorized") ) monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", mock_authenticate_user, ) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) client = TestClient(app) response = client.post( "/v2/login", json={"username": "alice", "password": "secret"}, ) assert response.status_code == 401 assert response.headers["content-type"] == "application/json" data = response.json() assert "error" in data assert isinstance(data["error"], dict) def test_login_v2_returns_json_on_unexpected_exception(monkeypatch): """Test that /v2/login returns JSON error when unexpected exception occurs""" mock_prisma_client = MagicMock() mock_authenticate_user = AsyncMock(side_effect=ValueError("Unexpected error")) monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", mock_authenticate_user, ) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) client = TestClient(app) response = client.post( "/v2/login", json={"username": "alice", "password": "secret"}, ) assert response.status_code == 500 assert response.headers["content-type"] == "application/json" data = response.json() assert "error" in data assert isinstance(data["error"], dict) assert "Unexpected error" in data["error"]["message"] def test_login_v2_returns_json_on_invalid_json_body(monkeypatch): """Test that /v2/login returns JSON error when request body is invalid JSON""" monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") client = TestClient(app) response = client.post( "/v2/login", content="invalid json", headers={"Content-Type": "application/json"}, ) assert response.status_code == 500 assert response.headers["content-type"] == "application/json" data = response.json() assert "error" in data assert isinstance(data["error"], dict) def test_login_v3_rejected_without_control_plane_url(monkeypatch): """v3/login returns 404 when control_plane_url is not configured.""" mock_prisma_client = MagicMock() monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) client = TestClient(app) response = client.post( "/v3/login", json={"username": "alice", "password": "secret"}, ) assert response.status_code == 404 assert "control_plane_url" in response.json()["error"]["message"] def test_login_v3_returns_code(monkeypatch): """v3/login returns an opaque code, not the JWT directly.""" mock_prisma_client = MagicMock() monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", AsyncMock(return_value={"user_id": "test-user"}), ) monkeypatch.setattr( "litellm.proxy.auth.login_utils.create_ui_token_object", MagicMock(return_value={"user_id": "test-user"}), ) monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr( "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"}, ) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_config = MagicMock() mock_config.worker_registry = [] monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) client = TestClient(app) response = client.post( "/v3/login", json={"username": "alice", "password": "secret"}, ) assert response.status_code == 200 data = response.json() assert "code" in data assert data["expires_in"] == 60 assert "token" not in data def test_login_v3_exchange_happy_path(monkeypatch): """Full flow: v3/login returns code, v3/login/exchange redeems it for JWT.""" mock_prisma_client = MagicMock() monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", AsyncMock(return_value={"user_id": "test-user"}), ) monkeypatch.setattr( "litellm.proxy.auth.login_utils.create_ui_token_object", MagicMock(return_value={"user_id": "test-user"}), ) monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr( "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"}, ) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_config = MagicMock() mock_config.worker_registry = [] monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) client = TestClient(app) # Step 1: login — get code login_response = client.post( "/v3/login", json={"username": "alice", "password": "secret"}, ) assert login_response.status_code == 200 code = login_response.json()["code"] # Step 2: exchange — get JWT exchange_response = client.post( "/v3/login/exchange", json={"code": code}, ) assert exchange_response.status_code == 200 exchange_data = exchange_response.json() assert exchange_data["token"] == "signed-token" assert "redirect_url" in exchange_data assert exchange_response.cookies.get("token") == "signed-token" def test_login_v3_exchange_single_use(monkeypatch): """Code can only be redeemed once.""" mock_prisma_client = MagicMock() monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", AsyncMock(return_value={"user_id": "test-user"}), ) monkeypatch.setattr( "litellm.proxy.auth.login_utils.create_ui_token_object", MagicMock(return_value={"user_id": "test-user"}), ) monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr( "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"}, ) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_config = MagicMock() mock_config.worker_registry = [] monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) client = TestClient(app) login_response = client.post( "/v3/login", json={"username": "alice", "password": "secret"}, ) code = login_response.json()["code"] # First exchange succeeds first = client.post("/v3/login/exchange", json={"code": code}) assert first.status_code == 200 # Second exchange fails second = client.post("/v3/login/exchange", json={"code": code}) assert second.status_code == 401 def test_login_v3_exchange_invalid_code(monkeypatch): """Random code returns 401.""" monkeypatch.setattr( "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"}, ) client = TestClient(app) response = client.post( "/v3/login/exchange", json={"code": "nonexistent-code"}, ) assert response.status_code == 401 def test_login_v3_exchange_rejected_without_control_plane_url(monkeypatch): """v3/login/exchange returns 404 when control_plane_url is not configured.""" monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) client = TestClient(app) response = client.post( "/v3/login/exchange", json={"code": "some-code"}, ) assert response.status_code == 404 assert "control_plane_url" in response.json()["error"]["message"] def test_login_v3_returns_json_on_proxy_exception(monkeypatch): """Test that /v3/login returns JSON error when ProxyException is raised""" from litellm.proxy._types import ProxyErrorTypes, ProxyException mock_prisma_client = MagicMock() mock_authenticate_user = AsyncMock( side_effect=ProxyException( message="Invalid credentials", type=ProxyErrorTypes.auth_error, param="password", code=401, ) ) monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", mock_authenticate_user, ) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") monkeypatch.setattr( "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"}, ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) client = TestClient(app) response = client.post( "/v3/login", json={"username": "alice", "password": "wrong"}, ) assert response.status_code == 401 assert response.headers["content-type"] == "application/json" data = response.json() assert "error" in data assert data["error"]["message"] == "Invalid credentials" assert data["error"]["type"] == "auth_error" def test_fallback_login_has_no_deprecation_banner(client_no_auth): response = client_no_auth.get("/fallback/login") assert response.status_code == 200 html = response.text assert '