wehub-resource-sync
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
"""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"
|