ui test fix

This commit is contained in:
Ishaan Jaffer
2025-09-23 14:34:08 -07:00
parent f8c1f519c9
commit bb6fb445c0
@@ -1303,7 +1303,7 @@ class TestCLIKeyRegenerationFlow:
)
# Assert
mock_create.assert_called_once_with(new_key)
mock_create.assert_called_once_with(key=new_key, user_id=None)
assert result.status_code == 200
assert "Success" in result.body.decode()
@@ -1320,8 +1320,17 @@ class TestCLIKeyRegenerationFlow:
# CLI state
cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-session-key-456"
# Mock the CLI callback
with patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback:
# Mock the CLI callback and required proxy server components
mock_result = {"user_id": "test-user", "email": "test@example.com"}
with patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback, \
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), \
patch("litellm.proxy.proxy_server.master_key", "test-master-key"), \
patch("litellm.proxy.proxy_server.general_settings", {}), \
patch("litellm.proxy.proxy_server.jwt_handler", MagicMock()), \
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \
patch.dict(os.environ, {"GOOGLE_CLIENT_ID": "test-google-id"}, clear=True), \
patch("litellm.proxy.management_endpoints.ui_sso.GoogleSSOHandler.get_google_callback_response", return_value=mock_result):
mock_cli_callback.return_value = MagicMock()
# Act
@@ -1329,9 +1338,10 @@ class TestCLIKeyRegenerationFlow:
# Assert
mock_cli_callback.assert_called_once_with(
mock_request,
request=mock_request,
key="sk-new-session-key-456",
existing_key="sk-existing-cli-key-123"
existing_key="sk-existing-cli-key-123",
result=mock_result
)
def test_get_redirect_url_preserves_existing_key(self):