chore: import upstream snapshot with attribution

tests/tracking/context/test_databricks_command_context.py

@@ -0,0 +1,21 @@
+from unittest import mock
+
+from mlflow.tracking.context.databricks_command_context import DatabricksCommandRunContext
+from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID
+
+
+def test_databricks_command_run_context_in_context():
+    with mock.patch("mlflow.utils.databricks_utils.get_job_group_id", return_value="1"):
+        assert DatabricksCommandRunContext().in_context()
+
+
+def test_databricks_command_run_context_tags():
+    with mock.patch("mlflow.utils.databricks_utils.get_job_group_id") as job_group_id_mock:
+        assert DatabricksCommandRunContext().tags() == {
+            MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID: job_group_id_mock.return_value
+        }
+
+
+def test_databricks_command_run_context_tags_nones():
+    with mock.patch("mlflow.utils.databricks_utils.get_job_group_id", return_value=None):
+        assert DatabricksCommandRunContext().tags() == {}