chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/databricks_test_utils.py

@@ -0,0 +1,69 @@
+"""Test utilities for Databricks datasource tests."""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from ray.data._internal.datasource.databricks_credentials import (
+    DatabricksCredentialProvider,
+)
+
+
+@dataclass
+class MockResponse:
+    """Mock HTTP response for testing.
+
+    Args:
+        status_code: HTTP status code. Defaults to 200.
+        content: Response content as bytes. Defaults to None.
+        _json_data: JSON response data. Defaults to None.
+        raise_on_error: If True, raise_for_status() raises for status >= 400.
+            Defaults to True.
+    """
+
+    status_code: int = 200
+    content: Optional[bytes] = None
+    _json_data: Optional[dict] = None
+    raise_on_error: bool = field(default=True, repr=False)
+
+    def raise_for_status(self):
+        """Raise an exception if status code indicates an error."""
+        if self.raise_on_error and self.status_code >= 400:
+            raise Exception(f"HTTP Error {self.status_code}")
+
+    def json(self):
+        """Return the JSON data."""
+        return self._json_data
+
+
+class RefreshableCredentialProvider(DatabricksCredentialProvider):
+    """A credential provider that simulates token refresh on invalidate.
+
+    Useful for testing 401 retry logic. When invalidate() is called,
+    the token changes from initial_token to "refreshed_token".
+
+    Args:
+        initial_token: The initial token value. Defaults to "expired_token".
+        host: The host URL to return. Defaults to "https://test-host.databricks.com".
+    """
+
+    def __init__(
+        self,
+        initial_token: str = "expired_token",
+        host: str = "https://test-host.databricks.com",
+    ):
+        self.current_token = initial_token
+        self.invalidate_count = 0
+        self._host = host
+
+    def get_token(self) -> str:
+        """Get the current token."""
+        return self.current_token
+
+    def get_host(self) -> str:
+        """Get the host URL."""
+        return self._host
+
+    def invalidate(self) -> None:
+        """Simulate token refresh by changing to 'refreshed_token'."""
+        self.invalidate_count += 1
+        self.current_token = "refreshed_token"