1  from unittest import mock
 2  
 3  import pytest
 4  
 5  from mlflow.tracing.utils.environment import resolve_env_metadata
 6  from mlflow.utils.mlflow_tags import (
 7      MLFLOW_DATABRICKS_NOTEBOOK_ID,
 8      MLFLOW_DATABRICKS_NOTEBOOK_PATH,
 9      MLFLOW_GIT_BRANCH,
10      MLFLOW_GIT_COMMIT,
11      MLFLOW_GIT_REPO_URL,
12      MLFLOW_SOURCE_NAME,
13      MLFLOW_SOURCE_TYPE,
14      MLFLOW_USER,
15  )
16  from mlflow.version import IS_TRACING_SDK_ONLY
17  
18  
19  @pytest.fixture(autouse=True)
20  def clear_lru_cache():
21      resolve_env_metadata.cache_clear()
22  
23  
24  def test_resolve_env_metadata():
25      expected_metadata = {
26          MLFLOW_USER: mock.ANY,
27          MLFLOW_SOURCE_NAME: mock.ANY,
28          MLFLOW_SOURCE_TYPE: "LOCAL",
29      }
30      if not IS_TRACING_SDK_ONLY:
31          expected_metadata.update({
32              MLFLOW_GIT_BRANCH: mock.ANY,
33              MLFLOW_GIT_COMMIT: mock.ANY,
34              MLFLOW_GIT_REPO_URL: mock.ANY,
35          })
36      assert resolve_env_metadata() == expected_metadata
37  
38  
39  def test_resolve_env_metadata_in_databricks_notebook():
40      with (
41          mock.patch(
42              "mlflow.tracking.context.databricks_notebook_context.databricks_utils"
43          ) as mock_db_utils,
44          mock.patch("mlflow.tracing.utils.environment.is_in_databricks_notebook", return_value=True),
45      ):
46          mock_db_utils.is_in_databricks_notebook.return_value = True
47          mock_db_utils.get_notebook_id.return_value = "notebook_123"
48          mock_db_utils.get_notebook_path.return_value = "/Users/bob/test.py"
49          mock_db_utils.get_webapp_url.return_value = None
50          mock_db_utils.get_workspace_url.return_value = None
51          mock_db_utils.get_workspace_id.return_value = None
52          mock_db_utils.get_workspace_info_from_dbutils.return_value = (None, None)
53  
54          assert resolve_env_metadata() == {
55              MLFLOW_USER: mock.ANY,
56              MLFLOW_SOURCE_NAME: "/Users/bob/test.py",
57              MLFLOW_SOURCE_TYPE: "NOTEBOOK",
58              MLFLOW_DATABRICKS_NOTEBOOK_ID: "notebook_123",
59              MLFLOW_DATABRICKS_NOTEBOOK_PATH: "/Users/bob/test.py",
60          }