wehub-resource-sync
51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.utils.rest_utils import MlflowHostCreds
|
|
|
|
_DATABRICKS_UC_REGISTRY_URI = "databricks-uc"
|
|
_DATABRICKS_TRACKING_URI = "databricks"
|
|
_DATABRICKS_UC_OSS_REGISTRY_URI = "uc"
|
|
_REGISTRY_HOST_CREDS = MlflowHostCreds("https://hello-registry")
|
|
_TRACKING_HOST_CREDS = MlflowHostCreds("https://hello-tracking")
|
|
|
|
|
|
def mock_host_creds(uri):
|
|
if uri == _DATABRICKS_TRACKING_URI:
|
|
return _TRACKING_HOST_CREDS
|
|
elif uri in (_DATABRICKS_UC_REGISTRY_URI, _DATABRICKS_UC_OSS_REGISTRY_URI):
|
|
return _REGISTRY_HOST_CREDS
|
|
raise Exception(f"Got unexpected store URI {uri}")
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_databricks_uc_host_creds():
|
|
with mock.patch(
|
|
"mlflow.store._unity_catalog.registry.rest_store.get_databricks_host_creds",
|
|
side_effect=mock_host_creds,
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_databricks_uc_oss_host_creds():
|
|
with mock.patch(
|
|
"mlflow.store._unity_catalog.registry.uc_oss_rest_store.get_oss_host_creds",
|
|
side_effect=mock_host_creds,
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def configure_client_for_uc(mock_databricks_uc_host_creds):
|
|
"""
|
|
Configure MLflow client to register models to UC
|
|
"""
|
|
with mock.patch("mlflow.utils.databricks_utils.get_databricks_host_creds"):
|
|
orig_registry_uri = mlflow.get_registry_uri()
|
|
mlflow.set_registry_uri("databricks-uc")
|
|
yield
|
|
mlflow.set_registry_uri(orig_registry_uri)
|