Python Test Helper for CodeChef

codechef_python_test_helper.py

import builtins
import sys
import textwrap

import mock


def mock_readline_and_assert_stdout(function_to_run, mocked_stdin, expected_stdout):
    #  Utility function that allow you to assert a given STDIN produce the expected STDOUT.
    #  Useful for TDDing CodeChef challenges: https://www.codechef.com/SNCKQL19/problems/QUALPREL
    #
    #  Note: This utility only works when accessing STDIN via the function `sys.stdin.readline`
    #
    #  Example use:
    #
    #      mock_readline_and_assert_stdout(
    #          function_to_run=FakeFunctions.read_and_echo_all_lines,
    #          mocked_stdin="""
    #              bonsoir
    #              eliot
    #          """,
    #          expected_stdout="""
    #              Echoed: bonsoir
    #              Echoed: eliot
    #          """)  # https://www.youtube.com/watch?v=BMVilMcSreU

    with mock.patch.object(builtins, 'print') as mocked_print:
        with mock.patch.object(sys.stdin, 'readline') as mocked_stdin_readline:
            def extract_lines(multiline_string):
                lines = textwrap.dedent(multiline_string).split('\n')
                if len(lines) > 2:
                    lines.pop(0)
                    lines.pop()
                return lines

            def inject_mocked_stdin_via_readline():
                mocked_stdin_readline.side_effect = extract_lines(mocked_stdin)

            def run_the_function_under_test():
                function_to_run()

            def assert_expected_stdout_was_printed():
                expected_stdout_lines = extract_lines(expected_stdout)

                def expected_calls():
                    return list(map(lambda line: mock.call(line), expected_stdout_lines))

                def indent(multiline_string):
                    number_of_spaces = 20
                    prefix = ''.join([' ' for i in range(0, number_of_spaces)])
                    return textwrap.indent(multiline_string, prefix)

                def build_stdout():
                    printed_lines = []
                    for call in mocked_print.call_args_list:
                        printed_line = call[0][0]
                        printed_lines.append(printed_line)

                    return indent('\n' + '\n'.join(printed_lines))

                formatted_expected_stdout = indent(textwrap.dedent(expected_stdout))
                assert mocked_print.call_args_list == expected_calls(), textwrap.dedent(f"""
                Expected STDOUT: {formatted_expected_stdout}
                Found STDOUT: {build_stdout()}
                """)

            inject_mocked_stdin_via_readline()
            run_the_function_under_test()
            assert_expected_stdout_was_printed()