from test_utils import *


def test_function_def():

    @Commentor("<return>")
    def target():
        pass

    asserteq_or_print(target(), '''
    def target():
        pass
''')


def test_return():

    with closing(StringIO()) as f:

        @Commentor(f)
        def target():
            a = 1
            return a + 1

        assert target() == 2

        asserteq_or_print(
            f.getvalue(), '''
        def target():
            a = 1
            return a + 1
            """
            1 : a
            2 : a + 1
            """
''')