ray-project--ray/python/ray/tests/test_submission_client_auth.py

import pytest

from ray._private.authentication.authentication_constants import (
    TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE,
    TOKEN_INVALID_ERROR_MESSAGE,
)
from ray._private.authentication_test_utils import (
    clear_auth_token_sources,
    reset_auth_token_state,
    set_auth_mode,
    set_env_auth_token,
)
from ray.dashboard.modules.dashboard_sdk import SubmissionClient
from ray.dashboard.modules.job.sdk import JobSubmissionClient
from ray.exceptions import AuthenticationError
from ray.util.state import StateApiClient


def test_submission_client_adds_token_automatically(setup_cluster_with_token_auth):
    """Test that SubmissionClient automatically adds token to headers."""
    # Token is already set in environment from setup_cluster_with_token_auth fixture

    client = SubmissionClient(address=setup_cluster_with_token_auth["dashboard_url"])

    # Verify authorization header was added (lowercase as per implementation)
    assert "authorization" in client._headers
    assert client._headers["authorization"].startswith("Bearer ")


def test_submission_client_without_token_shows_helpful_error(
    setup_cluster_with_token_auth,
):
    """Test that requests without token show helpful error message."""
    # Remove token from environment
    clear_auth_token_sources(remove_default=True)
    set_auth_mode("disabled")
    reset_auth_token_state()

    client = SubmissionClient(address=setup_cluster_with_token_auth["dashboard_url"])

    # Make a request - should fail with AuthenticationError
    with pytest.raises(AuthenticationError) as exc_info:
        client.get_version()

    error_str = str(exc_info.value)
    # Check that the error contains the simple message and auto-added docs link
    assert TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE in error_str


def test_submission_client_with_invalid_token_shows_helpful_error(
    setup_cluster_with_token_auth,
):
    """Test that requests with wrong token show helpful error message."""
    # Set wrong token
    wrong_token = "wrong_token_00000000000000000000000000000000"
    set_env_auth_token(wrong_token)
    set_auth_mode("token")
    reset_auth_token_state()

    client = SubmissionClient(address=setup_cluster_with_token_auth["dashboard_url"])

    # Make a request - should fail with AuthenticationError
    with pytest.raises(AuthenticationError) as exc_info:
        client.get_version()

    error_str = str(exc_info.value)
    # Check that the error contains the simple message and auto-added docs link
    assert TOKEN_INVALID_ERROR_MESSAGE in error_str


def test_submission_client_with_valid_token_succeeds(setup_cluster_with_token_auth):
    """Test that requests with valid token succeed."""
    client = SubmissionClient(address=setup_cluster_with_token_auth["dashboard_url"])

    # Make a request - should succeed
    version = client.get_version()
    assert version is not None


def test_job_submission_client_inherits_auth(setup_cluster_with_token_auth):
    """Test that JobSubmissionClient inherits auth from SubmissionClient."""
    client = JobSubmissionClient(address=setup_cluster_with_token_auth["dashboard_url"])

    # Verify authorization header was added (lowercase as per implementation)
    assert "authorization" in client._headers
    assert client._headers["authorization"].startswith("Bearer ")

    # Verify client can make authenticated requests
    version = client.get_version()
    assert version is not None


def test_state_api_client_inherits_auth(setup_cluster_with_token_auth):
    """Test that StateApiClient inherits auth from SubmissionClient."""
    client = StateApiClient(address=setup_cluster_with_token_auth["dashboard_url"])

    # Verify authorization header was added (lowercase as per implementation)
    assert "authorization" in client._headers
    assert client._headers["authorization"].startswith("Bearer ")


def test_user_provided_header_not_overridden(setup_cluster_with_token_auth):
    """Test that user-provided Authorization header is not overridden."""
    custom_auth = "Bearer custom_token"

    client = SubmissionClient(
        address=setup_cluster_with_token_auth["dashboard_url"],
        headers={"Authorization": custom_auth},
    )

    # Verify custom value is preserved
    assert client._headers["Authorization"] == custom_auth


def test_user_provided_header_case_insensitive(setup_cluster_with_token_auth):
    """Test that user-provided Authorization header is preserved regardless of case."""
    custom_auth = "Bearer custom_token"

    # Test with lowercase "authorization"
    client_lowercase = SubmissionClient(
        address=setup_cluster_with_token_auth["dashboard_url"],
        headers={"authorization": custom_auth},
    )

    # Verify custom value is preserved and no duplicate header added
    assert client_lowercase._headers["authorization"] == custom_auth
    assert "Authorization" not in client_lowercase._headers

    # Test with mixed case "AuThOrIzAtIoN"
    client_mixedcase = SubmissionClient(
        address=setup_cluster_with_token_auth["dashboard_url"],
        headers={"AuThOrIzAtIoN": custom_auth},
    )

    # Verify custom value is preserved and no duplicate header added
    assert client_mixedcase._headers["AuThOrIzAtIoN"] == custom_auth
    assert "Authorization" not in client_mixedcase._headers
    assert "authorization" not in client_mixedcase._headers


def test_error_messages_contain_instructions(setup_cluster_with_token_auth):
    """Test that all auth error messages contain setup instructions."""
    # Test 401 error (missing token)
    clear_auth_token_sources(remove_default=True)
    set_auth_mode("disabled")
    reset_auth_token_state()

    client = SubmissionClient(address=setup_cluster_with_token_auth["dashboard_url"])

    with pytest.raises(AuthenticationError) as exc_info:
        client.get_version()

    error_str = str(exc_info.value)
    assert TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE in error_str

    # Test 403 error (invalid token)
    set_env_auth_token("wrong_token_00000000000000000000000000000000")
    set_auth_mode("token")
    reset_auth_token_state()

    client2 = SubmissionClient(address=setup_cluster_with_token_auth["dashboard_url"])

    with pytest.raises(AuthenticationError) as exc_info:
        client2.get_version()

    error_str = str(exc_info.value)
    assert TOKEN_INVALID_ERROR_MESSAGE in error_str


def test_no_token_added_when_auth_disabled(setup_cluster_without_token_auth):
    """Test that no authorization header is injected when auth is disabled."""

    client = SubmissionClient(address=setup_cluster_without_token_auth["dashboard_url"])

    # Check both lowercase and uppercase variants
    assert "authorization" not in client._headers
    assert "Authorization" not in client._headers


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-vv", __file__]))