def test__get_block_variables_is_str(self):
pipeline = MLPipeline(['a_primitive'])
pipeline.blocks['a_primitive#1'].produce_outputs = 'get_produce_outputs'
pipeline.blocks['a_primitive#1'].instance.get_produce_outputs.return_value = [
{
'name': 'output_from_function',
'type': 'test'
}
]
outputs = pipeline._get_block_variables(
'a_primitive#1',
'produce_outputs',
{'output': 'name_output'}
)
expected = {
'output_from_function': {
'name': 'output_from_function',
'type': 'test',
}
}
assert outputs == expected
pipeline.blocks['a_primitive#1'].instance.get_produce_outputs.assert_called_once_with()
def test__get_block_variables_is_dict(self):
pipeline = MLPipeline(['a_primitive'])
pipeline.blocks['a_primitive#1'].produce_outputs = [{
'name': 'output',
'type': 'whatever'
}]
outputs = pipeline._get_block_variables('a_primitive#1',
'produce_outputs',
{'output': 'name_output'})
expected = {
'name_output': {
'name': 'output',
'type': 'whatever',
}
}
assert outputs == expected
