chore: import upstream snapshot with attribution

tests/tracking/request_header/test_databricks_request_header_provider.py

@@ -0,0 +1,91 @@
+import itertools
+from unittest import mock
+
+import pytest
+
+from mlflow.tracking.request_header.databricks_request_header_provider import (
+    DatabricksRequestHeaderProvider,
+)
+
+bool_values = [True, False]
+
+
+@pytest.mark.parametrize(
+    ("is_in_databricks_notebook", "is_in_databricks_job", "is_in_cluster"),
+    itertools.product(bool_values, bool_values, bool_values),
+)
+def test_databricks_request_header_provider_in_context(
+    is_in_databricks_notebook, is_in_databricks_job, is_in_cluster
+):
+    with (
+        mock.patch(
+            "mlflow.utils.databricks_utils.is_in_databricks_notebook",
+            return_value=is_in_databricks_notebook,
+        ),
+        mock.patch(
+            "mlflow.utils.databricks_utils.is_in_databricks_job", return_value=is_in_databricks_job
+        ),
+        mock.patch("mlflow.utils.databricks_utils.is_in_cluster", return_value=is_in_cluster),
+    ):
+        assert (
+            DatabricksRequestHeaderProvider().in_context() == is_in_databricks_notebook
+            or is_in_databricks_job
+            or is_in_cluster
+        )
+
+
+# test that request_headers returns whatever is available
+@pytest.mark.parametrize(
+    ("is_in_databricks_notebook", "is_in_databricks_job", "is_in_cluster"),
+    itertools.product(bool_values, bool_values, bool_values),
+)
+def test_databricks_request_header_provider_request_headers(
+    is_in_databricks_notebook, is_in_databricks_job, is_in_cluster
+):
+    with (
+        mock.patch(
+            "mlflow.utils.databricks_utils.is_in_databricks_notebook",
+            return_value=is_in_databricks_notebook,
+        ),
+        mock.patch(
+            "mlflow.utils.databricks_utils.is_in_databricks_job", return_value=is_in_databricks_job
+        ),
+        mock.patch("mlflow.utils.databricks_utils.is_in_cluster", return_value=is_in_cluster),
+        mock.patch("mlflow.utils.databricks_utils.get_notebook_id") as notebook_id_mock,
+        mock.patch("mlflow.utils.databricks_utils.get_job_id") as job_id_mock,
+        mock.patch("mlflow.utils.databricks_utils.get_job_run_id") as job_run_id_mock,
+        mock.patch("mlflow.utils.databricks_utils.get_job_type") as job_type_mock,
+        mock.patch("mlflow.utils.databricks_utils.get_cluster_id") as cluster_id_mock,
+        mock.patch("mlflow.utils.databricks_utils.get_workload_id") as workload_id_mock,
+        mock.patch("mlflow.utils.databricks_utils.get_workload_class") as workload_class_mock,
+    ):
+        request_headers = DatabricksRequestHeaderProvider().request_headers()
+
+        if is_in_databricks_notebook:
+            assert request_headers["notebook_id"] == notebook_id_mock.return_value
+        else:
+            assert "notebook_id" not in request_headers
+
+        if is_in_databricks_job:
+            assert request_headers["job_id"] == job_id_mock.return_value
+            assert request_headers["job_run_id"] == job_run_id_mock.return_value
+            assert request_headers["job_type"] == job_type_mock.return_value
+        else:
+            assert "job_id" not in request_headers
+            assert "job_run_id" not in request_headers
+            assert "job_type" not in request_headers
+
+        if is_in_cluster:
+            assert request_headers["cluster_id"] == cluster_id_mock.return_value
+        else:
+            assert "cluster_id" not in request_headers
+
+        if workload_id_mock.return_value is not None:
+            assert request_headers["workload_id"] == workload_id_mock.return_value
+        else:
+            assert "workload_id" not in request_headers
+
+        if workload_class_mock.return_value is not None:
+            assert request_headers["workload_class"] == workload_class_mock.return_value
+        else:
+            assert "workload_class" not in request_headers