wehub-resource-sync
84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow import MlflowClient
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.tracking.default_experiment.databricks_notebook_experiment_provider import (
|
|
DatabricksNotebookExperimentProvider,
|
|
)
|
|
from mlflow.utils.mlflow_tags import MLFLOW_EXPERIMENT_SOURCE_ID, MLFLOW_EXPERIMENT_SOURCE_TYPE
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_resolved_notebook_experiment_id():
|
|
DatabricksNotebookExperimentProvider._resolve_notebook_experiment_id.cache_clear()
|
|
|
|
|
|
def test_databricks_notebook_default_experiment_in_context():
|
|
with mock.patch("mlflow.utils.databricks_utils.is_in_databricks_notebook") as in_notebook_mock:
|
|
assert DatabricksNotebookExperimentProvider().in_context() == in_notebook_mock.return_value
|
|
|
|
|
|
def test_databricks_notebook_default_experiment_id():
|
|
with (
|
|
mock.patch.object(
|
|
MlflowClient,
|
|
"create_experiment",
|
|
side_effect=MlflowException(
|
|
message="Error message", error_code=INVALID_PARAMETER_VALUE
|
|
),
|
|
),
|
|
mock.patch(
|
|
"mlflow.utils.databricks_utils.get_notebook_path",
|
|
return_value="path",
|
|
),
|
|
mock.patch("mlflow.utils.databricks_utils.get_notebook_id") as patch_notebook_id,
|
|
):
|
|
assert (
|
|
DatabricksNotebookExperimentProvider().get_experiment_id()
|
|
== patch_notebook_id.return_value
|
|
)
|
|
|
|
|
|
def test_databricks_repo_notebook_default_experiment_gets_id_by_request():
|
|
with (
|
|
mock.patch(
|
|
"mlflow.utils.databricks_utils.get_notebook_id",
|
|
return_value=1234,
|
|
),
|
|
mock.patch(
|
|
"mlflow.utils.databricks_utils.get_notebook_path",
|
|
return_value="/Repos/path",
|
|
),
|
|
mock.patch.object(
|
|
MlflowClient, "create_experiment", return_value="experiment_id"
|
|
) as create_experiment_mock,
|
|
):
|
|
DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id = None
|
|
returned_id = DatabricksNotebookExperimentProvider().get_experiment_id()
|
|
assert returned_id == "experiment_id"
|
|
tags = {MLFLOW_EXPERIMENT_SOURCE_TYPE: "REPO_NOTEBOOK", MLFLOW_EXPERIMENT_SOURCE_ID: 1234}
|
|
create_experiment_mock.assert_called_once_with("/Repos/path", None, tags)
|
|
|
|
|
|
def test_databricks_repo_notebook_default_experiment_uses_fallback_notebook_id():
|
|
with (
|
|
mock.patch(
|
|
"mlflow.utils.databricks_utils.get_notebook_id",
|
|
return_value=1234,
|
|
),
|
|
mock.patch(
|
|
"mlflow.utils.databricks_utils.get_notebook_path",
|
|
return_value="/Repos/path",
|
|
),
|
|
mock.patch.object(MlflowClient, "create_experiment") as create_experiment_mock,
|
|
):
|
|
DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id = None
|
|
create_experiment_mock.side_effect = MlflowException(
|
|
message="not enabled", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
returned_id = DatabricksNotebookExperimentProvider().get_experiment_id()
|
|
assert returned_id == 1234
|