chore: import upstream snapshot with attribution

python/ray/tests/test_token_auth_integration_2.py

@@ -0,0 +1,363 @@
+"""Integration tests for token-based authentication in Ray (part 2).
+
+CLI, token loader, metrics, and dashboard agent tests.
+"""
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+import ray
+import ray.dashboard.consts as dashboard_consts
+from ray._common.network_utils import build_address
+from ray._common.test_utils import (
+    PrometheusTimeseries,
+    fetch_prometheus_timeseries,
+    wait_for_condition,
+)
+from ray._private.test_utils import client_test_enabled
+
+try:
+    from ray._raylet import AuthenticationTokenLoader
+
+    _RAYLET_AVAILABLE = True
+except ImportError:
+    _RAYLET_AVAILABLE = False
+    AuthenticationTokenLoader = None
+
+from ray._private.authentication_test_utils import (
+    authentication_env_guard,
+    clear_auth_token_sources,
+    reset_auth_token_state,
+    set_auth_mode,
+    set_auth_token_path,
+    set_env_auth_token,
+)
+
+pytestmark = pytest.mark.skipif(
+    not _RAYLET_AVAILABLE,
+    reason="Authentication tests require ray._raylet (not available in minimal installs)",
+)
+
+
+@pytest.fixture(autouse=True)
+def _auto_clean_token_sources(clean_token_sources):
+    """Opt in to the shared clean_token_sources fixture for every test."""
+    yield
+
+
+def _get_dashboard_agent_address(cluster_info):
+    """Get the dashboard agent HTTP address from a running cluster."""
+    # Get agent address from internal KV
+    node_id = ray.nodes()[0]["NodeID"]
+    key = f"{dashboard_consts.DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX}{node_id}"
+    agent_addr = ray.experimental.internal_kv._internal_kv_get(
+        key, namespace=ray._private.ray_constants.KV_NAMESPACE_DASHBOARD
+    )
+    if agent_addr:
+        ip, http_port, grpc_port = json.loads(agent_addr)
+        return f"http://{ip}:{http_port}"
+    return None
+
+
+def _wait_and_get_dashboard_agent_address(cluster_info, timeout=30):
+    """Waits for the dashboard agent address to become available and returns it."""
+
+    def agent_address_is_available():
+        return _get_dashboard_agent_address(cluster_info) is not None
+
+    wait_for_condition(agent_address_is_available, timeout=timeout)
+    return _get_dashboard_agent_address(cluster_info)
+
+
+@pytest.mark.skipif(
+    client_test_enabled(),
+    reason="Uses subprocess ray CLI, not compatible with client mode",
+)
+@pytest.mark.parametrize("use_generate", [True, False])
+def test_get_auth_token_cli(use_generate):
+    """Test ray get-auth-token CLI command."""
+    test_token = "a" * 64
+
+    with authentication_env_guard():
+        if use_generate:
+            # Test --generate flag (no token set)
+            clear_auth_token_sources(remove_default=True)
+            args = ["ray", "get-auth-token", "--generate"]
+        else:
+            # Test with existing token from env var
+            set_env_auth_token(test_token)
+            reset_auth_token_state()
+            args = ["ray", "get-auth-token"]
+
+        env = os.environ.copy()
+        result = subprocess.run(
+            args,
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+
+        assert result.returncode == 0, (
+            f"ray get-auth-token should succeed. "
+            f"stdout: {result.stdout}, stderr: {result.stderr}"
+        )
+
+        # Verify token is printed to stdout
+        token = result.stdout.strip()
+        assert len(token) == 64, token
+        assert all(c in "0123456789abcdef" for c in token), "Token should be hex"
+
+        if not use_generate:
+            # When using env var, should get exact token back
+            assert token == test_token
+
+        # Verify logs went to stderr (if --generate was used)
+        if use_generate:
+            assert (
+                "generating new authentication token..." in result.stderr.lower()
+            ), "Should log generation to stderr"
+
+
+@pytest.mark.skipif(
+    client_test_enabled(),
+    reason="Uses subprocess ray CLI, not compatible with client mode",
+)
+def test_get_auth_token_cli_no_token_no_generate():
+    """Test ray get-auth-token fails without token and without --generate."""
+    with authentication_env_guard():
+        reset_auth_token_state()
+        clear_auth_token_sources(remove_default=True)
+        env = os.environ.copy()
+
+        result = subprocess.run(
+            ["ray", "get-auth-token"],
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+
+        assert result.returncode != 0, "Should fail when no token and no --generate"
+        assert "error" in result.stderr.lower(), "Should print error to stderr"
+        assert "no" in result.stderr.lower() and "token" in result.stderr.lower()
+
+
+@pytest.mark.skipif(
+    client_test_enabled(),
+    reason="Uses subprocess ray CLI, not compatible with client mode",
+)
+def test_get_auth_token_cli_piping():
+    """Test that ray get-auth-token output can be piped."""
+    test_token = "b" * 64
+
+    with authentication_env_guard():
+        set_env_auth_token(test_token)
+        reset_auth_token_state()
+        env = os.environ.copy()
+
+        # Test piping: use token in shell pipeline
+        result = subprocess.run(
+            "ray get-auth-token | wc -c",
+            shell=True,
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+
+        assert result.returncode == 0
+        char_count = int(result.stdout.strip())
+        assert char_count == 64, f"Expected 64 chars (no newline), got {char_count}"
+
+
+@pytest.mark.skipif(
+    client_test_enabled(),
+    reason="Tests AuthenticationTokenLoader directly, no benefit testing this in client mode",
+)
+def test_missing_token_file_raises_authentication_error():
+    """Test that RAY_AUTH_TOKEN_PATH pointing to missing file raises AuthenticationError."""
+    with authentication_env_guard():
+        # Clear first, then set up the specific test scenario
+        clear_auth_token_sources(remove_default=True)
+        set_auth_mode("token")
+        set_auth_token_path(None, "/nonexistent/path/to/token")
+        reset_auth_token_state()
+
+        token_loader = AuthenticationTokenLoader.instance()
+
+        with pytest.raises(ray.exceptions.AuthenticationError) as exc_info:
+            token_loader.has_token()
+
+        # Verify error message is informative
+        assert str(Path("/nonexistent/path/to/token")) in str(exc_info.value)
+        assert "RAY_AUTH_TOKEN_PATH" in str(exc_info.value)
+
+
+@pytest.mark.skipif(
+    client_test_enabled(),
+    reason="Tests AuthenticationTokenLoader directly, no benefit testing this in client mode",
+)
+def test_empty_token_file_raises_authentication_error(tmp_path):
+    """Test that RAY_AUTH_TOKEN_PATH pointing to empty file raises AuthenticationError."""
+    token_file = tmp_path / "empty_token_file.txt"
+    with authentication_env_guard():
+        # Clear first, then set up the specific test scenario
+        clear_auth_token_sources(remove_default=True)
+        set_auth_mode("token")
+        set_auth_token_path("", token_file)
+        reset_auth_token_state()
+
+        token_loader = AuthenticationTokenLoader.instance()
+
+        with pytest.raises(ray.exceptions.AuthenticationError) as exc_info:
+            token_loader.has_token()
+
+        assert "cannot be opened or is empty" in str(exc_info.value)
+        assert str(token_file) in str(exc_info.value)
+
+
+@pytest.mark.skipif(
+    client_test_enabled(),
+    reason="Tests AuthenticationTokenLoader directly, no benefit testing this in client mode",
+)
+def test_no_token_with_auth_enabled_returns_false():
+    """Test that has_token(ignore_auth_mode=True) returns False when no token exists.
+
+    This allows the caller (ensure_token_if_auth_enabled) to decide whether
+    to generate a new token or raise an error.
+    """
+    with authentication_env_guard():
+        set_auth_mode("token")
+        clear_auth_token_sources(remove_default=True)
+        reset_auth_token_state()
+
+        token_loader = AuthenticationTokenLoader.instance()
+
+        # has_token(ignore_auth_mode=True) should return False, not raise an exception
+        result = token_loader.has_token(ignore_auth_mode=True)
+        assert result is False
+
+
+@pytest.mark.skipif(
+    client_test_enabled(),
+    reason="no benefit testing this in client mode",
+)
+def test_opentelemetry_metrics_with_token_auth(setup_cluster_with_token_auth):
+    """Test that OpenTelemetry metrics are exported with token authentication.
+
+    This test verifies that the C++ OpenTelemetryMetricRecorder correctly includes
+    the authentication token in its gRPC metadata when exporting metrics to the
+    metrics agent. If the auth headers are missing or incorrect, the metrics agent
+    would reject the requests and metrics wouldn't be collected.
+    """
+
+    cluster_info = setup_cluster_with_token_auth
+    cluster = cluster_info["cluster"]
+
+    # Get the metrics export address from the head node
+    head_node = cluster.head_node
+    prom_addresses = [
+        build_address(head_node.node_ip_address, head_node.metrics_export_port)
+    ]
+
+    timeseries = PrometheusTimeseries()
+
+    def verify_metrics_collected():
+        """Verify that metrics are being exported successfully."""
+        fetch_prometheus_timeseries(prom_addresses, timeseries)
+        metric_names = list(timeseries.metric_descriptors.keys())
+
+        # Check for core Ray metrics that are always exported
+        # These metrics are exported via the C++ OpenTelemetry recorder
+        expected_metrics = [
+            "ray_node_cpu_utilization",
+            "ray_node_mem_used",
+            "ray_node_disk_usage",
+        ]
+
+        # At least some metrics should be present
+        return len(metric_names) > 0 and any(
+            any(expected in name for name in metric_names)
+            for expected in expected_metrics
+        )
+
+    # Wait for metrics to be collected
+    # If auth wasn't working, the metrics agent would reject the exports
+    # and we wouldn't see any metrics
+    wait_for_condition(verify_metrics_collected, retry_interval_ms=1000)
+
+
+@pytest.mark.parametrize(
+    "token_type,expected_status",
+    [
+        ("none", 401),  # No token -> Unauthorized
+        ("valid", "not_auth_error"),  # Valid token -> passes auth (may get 404)
+        ("invalid", 403),  # Invalid token -> Forbidden
+    ],
+    ids=["no_token", "valid_token", "invalid_token"],
+)
+def test_dashboard_agent_auth(
+    token_type, expected_status, setup_cluster_with_token_auth
+):
+    """Test dashboard agent authentication with various token scenarios."""
+    import requests
+
+    cluster_info = setup_cluster_with_token_auth
+
+    agent_address = _wait_and_get_dashboard_agent_address(cluster_info)
+
+    # Build headers based on token type
+    headers = {}
+    if token_type == "valid":
+        headers["Authorization"] = f"Bearer {cluster_info['token']}"
+    elif token_type == "invalid":
+        headers["Authorization"] = "Bearer invalid_token_12345678901234567890"
+    # token_type == "none" -> no Authorization header
+
+    response = requests.get(
+        f"{agent_address}/api/job_agent/jobs/nonexistent/logs",
+        headers=headers,
+        timeout=5,
+    )
+
+    if expected_status == "not_auth_error":
+        # Valid token should pass auth (may get 404 for nonexistent job)
+        assert response.status_code not in (401, 403), (
+            f"Valid token should be accepted, got {response.status_code}: "
+            f"{response.text}"
+        )
+    else:
+        assert (
+            response.status_code == expected_status
+        ), f"Expected {expected_status}, got {response.status_code}: {response.text}"
+
+
+@pytest.mark.parametrize(
+    "endpoint",
+    ["/api/healthz", "/api/local_raylet_healthz"],
+    ids=["healthz", "local_raylet_healthz"],
+)
+def test_dashboard_agent_health_check_public(endpoint, setup_cluster_with_token_auth):
+    """Test that agent health check endpoints remain public without auth."""
+    import requests
+
+    cluster_info = setup_cluster_with_token_auth
+
+    agent_address = _wait_and_get_dashboard_agent_address(cluster_info)
+
+    # Health check endpoints should be accessible without auth
+    response = requests.get(f"{agent_address}{endpoint}", timeout=5)
+    assert response.status_code == 200, (
+        f"Health check {endpoint} should return 200 without auth, "
+        f"got {response.status_code}: {response.text}"
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-vv", __file__]))