def test_wf_container_task_multiple():
    square = ContainerTask(name='square', input_data_dir='/var/inputs', output_data_dir='/var/outputs', inputs=kwtypes(val=int), outputs=kwtypes(out=int), image='alpine', command=['sh', '-c', 'echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out'])
    sum = ContainerTask(name='sum', input_data_dir='/var/flyte/inputs', output_data_dir='/var/flyte/outputs', inputs=kwtypes(x=int, y=int), outputs=kwtypes(out=int), image='alpine', command=['sh', '-c', 'echo $(( {{.Inputs.x}} + {{.Inputs.y}} )) | tee /var/flyte/outputs/out'])

    @workflow
    def raw_container_wf(val1: int, val2: int) -> int:
        return sum(x=square(val=val1), y=square(val=val2))
    with task_mock(square) as square_mock, task_mock(sum) as sum_mock:
        square_mock.side_effect = (lambda val: (val * val))
        assert (square(val=10) == 100)
        sum_mock.side_effect = (lambda x, y: (x + y))
        assert (sum(x=10, y=10) == 20)
        assert (raw_container_wf(val1=10, val2=10) == 200)
def test_wf_container_task_multiple():
    square = ContainerTask(name='square', input_data_dir='/var/inputs', output_data_dir='/var/outputs', inputs=kwtypes(val=int), outputs=kwtypes(out=int), image='alpine', command=['sh', '-c', 'echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out'])
    sum = ContainerTask(name='sum', input_data_dir='/var/flyte/inputs', output_data_dir='/var/flyte/outputs', inputs=kwtypes(x=int, y=int), outputs=kwtypes(out=int), image='alpine', command=['sh', '-c', 'echo $(( {{.Inputs.x}} + {{.Inputs.y}} )) | tee /var/flyte/outputs/out'])

    @workflow
    def raw_container_wf(val1: int, val2: int) -> int:
        return sum(x=square(val=val1), y=square(val=val2))
    with task_mock(square) as square_mock, task_mock(sum) as sum_mock:
        square_mock.side_effect = (lambda val: (val * val))
        assert (square(val=10) == 100)
        sum_mock.side_effect = (lambda x, y: (x + y))
        assert (sum(x=10, y=10) == 20)
        assert (raw_container_wf(val1=10, val2=10) == 200)