wehub-resource-sync
92 lines
3.7 KiB
Python
92 lines
3.7 KiB
Python
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
|