def test_get_store_for_unregistered_scheme():
tracking_store = TrackingStoreRegistry()
with pytest.raises(mlflow.exceptions.MlflowException,
match="Unexpected URI scheme"):
tracking_store.get_store("unknown-scheme://")
def test_get_store_for_unregistered_scheme():
tracking_store = TrackingStoreRegistry()
with pytest.raises(
UnsupportedModelRegistryStoreURIException,
match="Model registry functionality is unavailable",
):
tracking_store.get_store("unknown-scheme://")
def test_plugin_registration():
tracking_store = TrackingStoreRegistry()
test_uri = "mock-scheme://fake-host/fake-path"
test_scheme = "mock-scheme"
mock_plugin = mock.Mock()
tracking_store.register(test_scheme, mock_plugin)
assert test_scheme in tracking_store._registry
assert tracking_store.get_store(test_uri) == mock_plugin.return_value
mock_plugin.assert_called_once_with(store_uri=test_uri, artifact_uri=None)
def test_plugin_registration_via_entrypoints():
mock_plugin_function = mock.Mock()
mock_entrypoint = mock.Mock(load=mock.Mock(return_value=mock_plugin_function))
mock_entrypoint.name = "mock-scheme"
with mock.patch(
"entrypoints.get_group_all", return_value=[mock_entrypoint]
) as mock_get_group_all:
tracking_store = TrackingStoreRegistry()
tracking_store.register_entrypoints()
assert tracking_store.get_store("mock-scheme://") == mock_plugin_function.return_value
mock_plugin_function.assert_called_once_with(store_uri="mock-scheme://", artifact_uri=None)
mock_get_group_all.assert_called_once_with("mlflow.tracking_store")
def test_get_store_for_unregistered_scheme():
tracking_store = TrackingStoreRegistry()
with pytest.raises(UnsupportedModelRegistryStoreURIException):
tracking_store.get_store("unknown-scheme://")