Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test__load_pipeline_value_error(load_mock, gpp_mock):
load_mock.return_value = None
gpp_mock.return_value = ['a', 'b']
with pytest.raises(ValueError):
discovery.load_pipeline('invalid.pipeline')
load_mock.assert_called_once_with('invalid.pipeline', ['a', 'b'])
def test__load_pipeline_success(load_mock, gpp_mock):
gpp_mock.return_value = ['a', 'b']
pipeline = discovery.load_pipeline('valid.pipeline')
load_mock.assert_called_once_with('valid.pipeline', ['a', 'b'])
assert pipeline == load_mock.return_value
def _get_pipeline_dict(pipeline, primitives):
if isinstance(pipeline, dict):
return pipeline
elif isinstance(pipeline, str):
return load_pipeline(pipeline)
elif isinstance(pipeline, MLPipeline):
return pipeline.to_dict()
elif isinstance(pipeline, list):
if primitives is not None:
raise ValueError('if `pipeline` is a `list`, `primitives` must be `None`')
return {'primitives': pipeline}
elif pipeline is None:
if primitives is None:
raise ValueError('Either `pipeline` or `primitives` must be not `None`.')
return dict()