wehub-resource-sync
130 lines
4.7 KiB
Python
130 lines
4.7 KiB
Python
import time
|
|
from datetime import timedelta
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from databricks.sdk.service.sql import State
|
|
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_SQL_WAREHOUSE_AUTO_START,
|
|
MLFLOW_SQL_WAREHOUSE_AUTO_START_TIMEOUT_SECONDS,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.utils import databricks_sql_warehouse
|
|
from mlflow.utils.databricks_sql_warehouse import ensure_sql_warehouse_running
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clear_cache():
|
|
databricks_sql_warehouse._verified_running.clear()
|
|
yield
|
|
databricks_sql_warehouse._verified_running.clear()
|
|
|
|
|
|
def _make_mock_client(state_value):
|
|
client = mock.MagicMock()
|
|
warehouse_info = mock.MagicMock()
|
|
warehouse_info.state = State[state_value]
|
|
client.warehouses.get.return_value = warehouse_info
|
|
return client
|
|
|
|
|
|
def _patch_client(client):
|
|
return mock.patch.object(databricks_sql_warehouse, "_get_workspace_client", return_value=client)
|
|
|
|
|
|
def test_running_warehouse_skips_start():
|
|
client = _make_mock_client("RUNNING")
|
|
with _patch_client(client):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
client.warehouses.get.assert_called_once_with("wh-1")
|
|
client.warehouses.start_and_wait.assert_not_called()
|
|
|
|
|
|
def test_stopped_warehouse_starts_and_waits_with_default_timeout():
|
|
client = _make_mock_client("STOPPED")
|
|
with (
|
|
_patch_client(client),
|
|
mock.patch.object(databricks_sql_warehouse._logger, "info") as log_info,
|
|
):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
client.warehouses.start_and_wait.assert_called_once_with(
|
|
"wh-1", timeout=timedelta(seconds=1200)
|
|
)
|
|
log_info.assert_called_once()
|
|
rendered = log_info.call_args.args[0]
|
|
assert "wh-1" in rendered
|
|
assert "STOPPED" in rendered
|
|
|
|
|
|
@pytest.mark.parametrize("state", ["STOPPING", "STARTING"])
|
|
def test_transitional_state_triggers_start_and_wait(state):
|
|
client = _make_mock_client(state)
|
|
with _patch_client(client):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
client.warehouses.start_and_wait.assert_called_once()
|
|
|
|
|
|
def test_second_call_within_ttl_is_cached(monkeypatch):
|
|
monkeypatch.setattr(databricks_sql_warehouse, "_CACHE_TTL_SECONDS", 60.0)
|
|
client = _make_mock_client("RUNNING")
|
|
with _patch_client(client):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
ensure_sql_warehouse_running("wh-1")
|
|
assert client.warehouses.get.call_count == 1
|
|
|
|
|
|
def test_second_call_after_ttl_rechecks(monkeypatch):
|
|
monkeypatch.setattr(databricks_sql_warehouse, "_CACHE_TTL_SECONDS", 0.0)
|
|
client = _make_mock_client("RUNNING")
|
|
with _patch_client(client):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
# TTL = 0, so the cache entry is expired immediately.
|
|
time.sleep(0.001)
|
|
ensure_sql_warehouse_running("wh-1")
|
|
assert client.warehouses.get.call_count == 2
|
|
|
|
|
|
def test_different_warehouse_ids_have_independent_cache():
|
|
client = _make_mock_client("RUNNING")
|
|
with _patch_client(client):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
ensure_sql_warehouse_running("wh-2")
|
|
assert client.warehouses.get.call_count == 2
|
|
client.warehouses.get.assert_any_call("wh-1")
|
|
client.warehouses.get.assert_any_call("wh-2")
|
|
|
|
|
|
def test_env_var_timeout_is_honored(monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_SQL_WAREHOUSE_AUTO_START_TIMEOUT_SECONDS.name, "30")
|
|
client = _make_mock_client("STOPPED")
|
|
with _patch_client(client):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
client.warehouses.start_and_wait.assert_called_once_with("wh-1", timeout=timedelta(seconds=30))
|
|
|
|
|
|
def test_timeout_error_raises_mlflow_exception():
|
|
client = _make_mock_client("STOPPED")
|
|
client.warehouses.start_and_wait.side_effect = TimeoutError("deadline exceeded")
|
|
with _patch_client(client), pytest.raises(MlflowException, match="wh-1") as exc_info:
|
|
ensure_sql_warehouse_running("wh-1")
|
|
assert MLFLOW_SQL_WAREHOUSE_AUTO_START_TIMEOUT_SECONDS.name in str(exc_info.value)
|
|
|
|
|
|
def test_generic_sdk_error_raises_mlflow_exception():
|
|
client = _make_mock_client("STOPPED")
|
|
client.warehouses.start_and_wait.side_effect = RuntimeError("warehouse not found")
|
|
with _patch_client(client), pytest.raises(MlflowException, match="wh-1") as exc_info:
|
|
ensure_sql_warehouse_running("wh-1")
|
|
assert "warehouse not found" in str(exc_info.value)
|
|
assert "MLFLOW_SQL_WAREHOUSE_AUTO_START" in str(exc_info.value)
|
|
|
|
|
|
def test_auto_start_disabled_is_noop(monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_SQL_WAREHOUSE_AUTO_START.name, "false")
|
|
client = _make_mock_client("STOPPED")
|
|
with _patch_client(client):
|
|
ensure_sql_warehouse_running("wh-1")
|
|
client.warehouses.get.assert_not_called()
|
|
client.warehouses.start_and_wait.assert_not_called()
|