wehub-resource-sync
52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.tracing.databricks import set_databricks_monitoring_sql_warehouse_id
|
|
|
|
|
|
def test_set_databricks_monitoring_sql_warehouse_id_requires_databricks_tracking_uri():
|
|
with mock.patch("mlflow.get_tracking_uri", return_value="file:///tmp"):
|
|
with pytest.raises(MlflowException, match="only supported when the tracking URI"):
|
|
set_databricks_monitoring_sql_warehouse_id(
|
|
sql_warehouse_id="warehouse123", experiment_id="exp456"
|
|
)
|
|
|
|
|
|
def test_set_databricks_monitoring_sql_warehouse_id_with_explicit_experiment_id():
|
|
mock_store = mock.MagicMock()
|
|
with (
|
|
mock.patch("mlflow.tracking.get_tracking_uri", return_value="databricks"),
|
|
mock.patch(
|
|
"mlflow.tracking._tracking_service.utils._get_store",
|
|
return_value=mock_store,
|
|
),
|
|
):
|
|
set_databricks_monitoring_sql_warehouse_id(
|
|
sql_warehouse_id="warehouse123", experiment_id="exp456"
|
|
)
|
|
mock_store.set_experiment_tag.assert_called_once()
|
|
call_args = mock_store.set_experiment_tag.call_args
|
|
assert call_args[0][0] == "exp456"
|
|
assert call_args[0][1].key == "mlflow.monitoring.sqlWarehouseId"
|
|
assert call_args[0][1].value == "warehouse123"
|
|
|
|
|
|
def test_set_databricks_monitoring_sql_warehouse_id_with_default_experiment_id():
|
|
mock_store = mock.MagicMock()
|
|
with (
|
|
mock.patch("mlflow.tracking.get_tracking_uri", return_value="databricks"),
|
|
mock.patch(
|
|
"mlflow.tracking._tracking_service.utils._get_store",
|
|
return_value=mock_store,
|
|
),
|
|
mock.patch("mlflow.tracking.fluent._get_experiment_id", return_value="default_exp"),
|
|
):
|
|
set_databricks_monitoring_sql_warehouse_id(sql_warehouse_id="warehouse789")
|
|
mock_store.set_experiment_tag.assert_called_once()
|
|
call_args = mock_store.set_experiment_tag.call_args
|
|
assert call_args[0][0] == "default_exp"
|
|
assert call_args[0][1].key == "mlflow.monitoring.sqlWarehouseId"
|
|
assert call_args[0][1].value == "warehouse789"
|