wehub-resource-sync
61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.tracing.utils.environment import resolve_env_metadata
|
|
from mlflow.utils.mlflow_tags import (
|
|
MLFLOW_DATABRICKS_NOTEBOOK_ID,
|
|
MLFLOW_DATABRICKS_NOTEBOOK_PATH,
|
|
MLFLOW_GIT_BRANCH,
|
|
MLFLOW_GIT_COMMIT,
|
|
MLFLOW_GIT_REPO_URL,
|
|
MLFLOW_SOURCE_NAME,
|
|
MLFLOW_SOURCE_TYPE,
|
|
MLFLOW_USER,
|
|
)
|
|
from mlflow.version import IS_TRACING_SDK_ONLY
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_lru_cache():
|
|
resolve_env_metadata.cache_clear()
|
|
|
|
|
|
def test_resolve_env_metadata():
|
|
expected_metadata = {
|
|
MLFLOW_USER: mock.ANY,
|
|
MLFLOW_SOURCE_NAME: mock.ANY,
|
|
MLFLOW_SOURCE_TYPE: "LOCAL",
|
|
}
|
|
if not IS_TRACING_SDK_ONLY:
|
|
expected_metadata.update({
|
|
MLFLOW_GIT_BRANCH: mock.ANY,
|
|
MLFLOW_GIT_COMMIT: mock.ANY,
|
|
MLFLOW_GIT_REPO_URL: mock.ANY,
|
|
})
|
|
assert resolve_env_metadata() == expected_metadata
|
|
|
|
|
|
def test_resolve_env_metadata_in_databricks_notebook():
|
|
with (
|
|
mock.patch(
|
|
"mlflow.tracking.context.databricks_notebook_context.databricks_utils"
|
|
) as mock_db_utils,
|
|
mock.patch("mlflow.tracing.utils.environment.is_in_databricks_notebook", return_value=True),
|
|
):
|
|
mock_db_utils.is_in_databricks_notebook.return_value = True
|
|
mock_db_utils.get_notebook_id.return_value = "notebook_123"
|
|
mock_db_utils.get_notebook_path.return_value = "/Users/bob/test.py"
|
|
mock_db_utils.get_webapp_url.return_value = None
|
|
mock_db_utils.get_workspace_url.return_value = None
|
|
mock_db_utils.get_workspace_id.return_value = None
|
|
mock_db_utils.get_workspace_info_from_dbutils.return_value = (None, None)
|
|
|
|
assert resolve_env_metadata() == {
|
|
MLFLOW_USER: mock.ANY,
|
|
MLFLOW_SOURCE_NAME: "/Users/bob/test.py",
|
|
MLFLOW_SOURCE_TYPE: "NOTEBOOK",
|
|
MLFLOW_DATABRICKS_NOTEBOOK_ID: "notebook_123",
|
|
MLFLOW_DATABRICKS_NOTEBOOK_PATH: "/Users/bob/test.py",
|
|
}
|