def test_get_diagram_multiple_blocks(self): f = open('tests/data/diagrams/diagram_multiple_blocks.txt', 'r') expected = f.read()[:-1] f.close() first_output = [ { 'name': 'output_variable_a', 'type': 'another_whatever', 'variable': 'a_primitive#1.output_variable_a' } ] second_output = [ { 'name': 'output_variable_b', 'type': 'another_whatever', 'variable': 'b_primitive#1.output_variable_b' } ] pipeline = MLPipeline(['a_primitive', 'b_primitive'], outputs={'default': second_output}) pipeline.blocks['a_primitive#1'].produce_args = [ { 'name': 'input_variable', 'type': 'whatever' } ] pipeline.blocks['a_primitive#1'].produce_output = first_output pipeline.blocks['b_primitive#1'].produce_args = first_output pipeline.blocks['b_primitive#1'].produce_output = second_output assert str(pipeline.get_diagram()) == expected
def test_get_diagram_fit(self): f = open('tests/data/diagrams/diagram_fit.txt', 'r') expected = f.read()[:-1] f.close() output = [ { 'name': 'output_variable', 'type': 'another_whatever', 'variable': 'a_primitive#1.output_variable' } ] pipeline = MLPipeline(['a_primitive'], outputs={'default': output}) pipeline.blocks['a_primitive#1'].produce_args = [ { 'name': 'input_variable', 'type': 'whatever' } ] pipeline.blocks['a_primitive#1'].fit_args = [ { 'name': 'input_variable', 'type': 'whatever' } ] pipeline.blocks['a_primitive#1'].produce_output = output assert str(pipeline.get_diagram()) == expected