def test_step_cloner_should_save_sub_steps(tmpdir):
    tape = TapeCallbackFunction()
    p = StepClonerForEachDataInput(Pipeline(
        [FitCallbackStep(tape), MultiplyByN(2)]),
                                   cache_folder_when_no_handle=tmpdir)
    data_inputs = _create_data((2, 2))
    expected_outputs = _create_data((2, 2))
    p, processed_outputs = p.fit_transform(data_inputs, expected_outputs)

    p.save(ExecutionContext(tmpdir), full_dump=True)

    saved_paths = [
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline[0]/FitCallbackStep/FitCallbackStep.joblib'
        ),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline[0]/MultiplyByN/MultiplyByN.joblib'
        ),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline[0]/MultiplyByN/MultiplyByN.joblib'
        ),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline[0]/Pipeline[0].joblib'),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline[1]/FitCallbackStep/FitCallbackStep.joblib'
        ),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline[1]/MultiplyByN/MultiplyByN.joblib'
        ),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline[1]/Pipeline[1].joblib'),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline/FitCallbackStep/FitCallbackStep.joblib'
        ),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/Pipeline/MultiplyByN/MultiplyByN.joblib'
        ),
        os.path.join(tmpdir,
                     'StepClonerForEachDataInput/Pipeline/Pipeline.joblib'),
        os.path.join(
            tmpdir,
            'StepClonerForEachDataInput/StepClonerForEachDataInput.joblib')
    ]

    for p in saved_paths:
        assert os.path.exists(p)
def test_step_cloner_should_load_sub_steps(tmpdir):
    tape = TapeCallbackFunction()
    p = StepClonerForEachDataInput(Pipeline(
        [FitCallbackStep(tape), MultiplyByN(2)]),
                                   cache_folder_when_no_handle=tmpdir)
    data_inputs = _create_data((2, 2))
    expected_outputs = _create_data((2, 2))
    p, processed_outputs = p.fit_transform(data_inputs, expected_outputs)

    p.save(ExecutionContext(tmpdir), full_dump=True)

    loaded_step_cloner = ExecutionContext(tmpdir).load(
        'StepClonerForEachDataInput')
    assert isinstance(loaded_step_cloner.wrapped, Pipeline)
    assert len(loaded_step_cloner.steps_as_tuple) == len(data_inputs)
    assert isinstance(loaded_step_cloner.steps_as_tuple[0][1], Pipeline)
    assert isinstance(loaded_step_cloner.steps_as_tuple[1][1], Pipeline)