From 6686a1213fdb09d97d0ae3b1931587f561fa0e0d Mon Sep 17 00:00:00 2001
From: Ryan Crabbe <rcrabbe@berkeley.edu>
Date: Sat, 14 Mar 2026 11:24:06 -0700
Subject: [PATCH] Fix login test fragility and add control-plane token test

Explicitly mock proxy_config.worker_registry in the existing login test
to isolate it from global state. Add positive test verifying the token
is returned in the response body when workers are configured.
---
 tests/test_litellm/proxy/test_proxy_server.py | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 864383f79ab..219a0ca7825 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -95,6 +95,9 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch):
     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)
+    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)
 
@@ -129,6 +132,43 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch):
     )
 
 
+def test_login_v2_includes_token_in_body_when_control_plane(monkeypatch):
+    """Token should appear in the response body when workers are configured."""
+    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", {})
+    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 = [
+        WorkerRegistryEntry(
+            worker_id="team-a", name="Team A", url="https://worker-1:4001"
+        ),
+    ]
+    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(
+        "/v2/login",
+        json={"username": "alice", "password": "secret"},
+    )
+
+    assert response.status_code == 200
+    assert response.json()["token"] == "signed-token"
+    assert response.cookies.get("token") == "signed-token"
+
+
 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