chore: import upstream snapshot with attribution

python/ray/_private/authentication_test_utils.py

@@ -0,0 +1,157 @@
+import os
+import shutil
+import tempfile
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Optional
+
+from ray._raylet import AuthenticationTokenLoader, Config
+
+_AUTH_ENV_VARS = ("RAY_AUTH_MODE", "RAY_AUTH_TOKEN", "RAY_AUTH_TOKEN_PATH")
+_DEFAULT_AUTH_TOKEN_RELATIVE_PATH = Path(".ray") / "auth_token"
+
+
+def reset_auth_token_state() -> None:
+    """Reset authentication token and AUTH_MODE ray config."""
+
+    AuthenticationTokenLoader.instance().reset_cache()
+    Config.initialize("")
+
+
+def set_auth_mode(mode: str) -> None:
+    """Set the authentication mode environment variable."""
+
+    os.environ["RAY_AUTH_MODE"] = mode
+
+
+def set_env_auth_token(token: str) -> None:
+    """Configure the authentication token via environment variable."""
+
+    os.environ["RAY_AUTH_TOKEN"] = token
+    os.environ.pop("RAY_AUTH_TOKEN_PATH", None)
+
+
+def set_auth_token_path(token: str, path: Path) -> None:
+    """Write the authentication token to a specific path and point the loader to it."""
+
+    token_path = Path(path)
+    if token is not None:
+        token_path.parent.mkdir(parents=True, exist_ok=True)
+        token_path.write_text(token)
+    os.environ["RAY_AUTH_TOKEN_PATH"] = str(token_path)
+    os.environ.pop("RAY_AUTH_TOKEN", None)
+
+
+def set_default_auth_token(token: str) -> Path:
+    """Write the authentication token to the default ~/.ray/auth_token location."""
+
+    default_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH
+    default_path.parent.mkdir(parents=True, exist_ok=True)
+    default_path.write_text(token)
+    return default_path
+
+
+def clear_auth_token_sources(remove_default: bool = False) -> None:
+    """Clear authentication-related environment variables and optional default token file."""
+
+    for var in ("RAY_AUTH_TOKEN", "RAY_AUTH_TOKEN_PATH"):
+        os.environ.pop(var, None)
+
+    if remove_default:
+        default_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH
+        default_path.unlink(missing_ok=True)
+
+
+@dataclass
+class AuthenticationEnvSnapshot:
+    original_env: Dict[str, Optional[str]]
+    original_home: Optional[str]
+    home_was_set: bool
+    temp_home: Optional[Path]
+    default_token_path: Path
+    default_token_exists: bool
+    default_token_contents: Optional[str]
+
+    @classmethod
+    def capture(cls) -> "AuthenticationEnvSnapshot":
+        """Capture current authentication-related environment state."""
+
+        original_env = {var: os.environ.get(var) for var in _AUTH_ENV_VARS}
+        home_was_set = "HOME" in os.environ
+        original_home = os.environ.get("HOME")
+        temp_home: Optional[Path] = None
+
+        if not home_was_set:
+            # in CI $HOME may not be set which can cause issues with tests related to default auth token file.
+            test_tmpdir = os.environ.get("TEST_TMPDIR")
+            base_dir = Path(test_tmpdir) if test_tmpdir else Path(tempfile.gettempdir())
+            temp_home = base_dir / "ray_test_home"
+            temp_home.mkdir(parents=True, exist_ok=True)
+            os.environ["HOME"] = str(temp_home)
+
+        default_token_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH
+        default_token_exists = default_token_path.exists()
+        default_token_contents = (
+            default_token_path.read_text() if default_token_exists else None
+        )
+
+        return cls(
+            original_env=original_env,
+            original_home=original_home,
+            home_was_set=home_was_set,
+            temp_home=temp_home,
+            default_token_path=default_token_path,
+            default_token_exists=default_token_exists,
+            default_token_contents=default_token_contents,
+        )
+
+    def clear_default_token(self) -> None:
+        """Remove the default token file for the current HOME."""
+
+        self.default_token_path.unlink(missing_ok=True)
+
+    def restore(self) -> None:
+        """Restore the captured environment, HOME, and default token file state."""
+        # delete any custom token files that may have been created during the test
+        custom_token_path = os.environ.get("RAY_AUTH_TOKEN_PATH")
+        if custom_token_path is not None:
+            custom_token_path = Path(custom_token_path)
+            if custom_token_path.exists():
+                custom_token_path.unlink(missing_ok=True)
+
+        for var, value in self.original_env.items():
+            if value is None:
+                os.environ.pop(var, None)
+            else:
+                os.environ[var] = value
+
+        if self.home_was_set:
+            if self.original_home is None:
+                os.environ.pop("HOME", None)
+            else:
+                os.environ["HOME"] = self.original_home
+
+        if self.default_token_exists:
+            self.default_token_path.parent.mkdir(parents=True, exist_ok=True)
+            self.default_token_path.write_text(self.default_token_contents or "")
+        else:
+            self.default_token_path.unlink(missing_ok=True)
+
+        if not self.home_was_set:
+            current_home = os.environ.get("HOME")
+            if self.temp_home is not None and current_home == str(self.temp_home):
+                os.environ.pop("HOME", None)
+            if self.temp_home is not None and self.temp_home.exists():
+                shutil.rmtree(self.temp_home, ignore_errors=True)
+
+
+@contextmanager
+def authentication_env_guard():
+    """Context manager that restores authentication environment state on exit."""
+
+    snapshot = AuthenticationEnvSnapshot.capture()
+    try:
+        yield snapshot
+    finally:
+        snapshot.restore()