1  from unittest import mock
 2  
 3  import pytest
 4  
 5  from mlflow.exceptions import MlflowException
 6  from mlflow.store.model_registry.rest_store import RestStore
 7  from mlflow.utils.rest_utils import MlflowHostCreds
 8  
 9  ACTIVE_WORKSPACE = "team-a"
10  
11  
12  def test_model_registry_rest_store_workspace_guard():
13      creds = MlflowHostCreds("https://hello")
14      store = RestStore(lambda: creds)
15      store._workspace_support = False
16  
17      with mock.patch(
18          "mlflow.store.workspace_rest_store_mixin.get_request_workspace",
19          return_value=ACTIVE_WORKSPACE,
20      ):
21          with pytest.raises(
22              MlflowException,
23              match="Active workspace 'team-a' cannot be used because the remote server does not",
24          ):
25              store.search_registered_models()