def test_add_function(self):

        code = "def foo(a, b):\n" "    x, y = a, b\n" "    return x, y\n" "c, d = foo(1, 2)"

        block = Block(code)
        info = find_functions(block.ast)
        func = FunctionCall.from_ast(block.sub_blocks[-1].ast, info)
        model = ExecutionModel()
        model.add_function(func)

        assert len(model.sorted_statements) == 1

        desired = "\ndef foo(a, b): \n    x, y = (a, b)\n    return x, y\n\nx_, y_ = foo(1, 2)"
        self.assertEqual(desired, model.code)
    def test_add_function(self):

        code = "def foo(a, b):\n" \
                "    x, y = a, b\n" \
                "    return x, y\n" \
                "c, d = foo(1, 2)"

        block = Block(code)
        info = find_functions(block.ast)
        func = FunctionCall.from_ast(block.sub_blocks[-1].ast, info)
        model = ExecutionModel()
        model.add_function(func)

        assert len(model.sorted_statements) == 1

        desired = '\ndef foo(a, b): \n    x, y = (a, b)\n    return x, y\n\nx_, y_ = foo(1, 2)'
        self.assertEqual(desired, model.code)
    def test_remove_function(self):

        code = "def foo(a, b):\n" "    x, y = a, b\n" "    return x, y\n" "c, d = foo(1, 2)\n" "e, f = foo(3, 4)"

        block = Block(code)
        info = find_functions(block.ast)
        foo1_func = FunctionCall.from_ast(block.sub_blocks[-1].ast, info)
        foo2_func = FunctionCall.from_ast(block.sub_blocks[-2].ast, info)
        assert foo1_func != foo2_func

        model = ExecutionModel()
        model.add_function(foo1_func)
        model.add_function(foo2_func)

        model.remove_function(foo2_func)

        assert not (foo2_func in model.statements)
        assert foo1_func in model.statements
    def test_remove_function(self):

        code = "def foo(a, b):\n" \
                "    x, y = a, b\n" \
                "    return x, y\n" \
                "c, d = foo(1, 2)\n" \
                "e, f = foo(3, 4)"

        block = Block(code)
        info = find_functions(block.ast)
        foo1_func = FunctionCall.from_ast(block.sub_blocks[-1].ast, info)
        foo2_func = FunctionCall.from_ast(block.sub_blocks[-2].ast, info)
        assert foo1_func != foo2_func

        model = ExecutionModel()
        model.add_function(foo1_func)
        model.add_function(foo2_func)

        model.remove_function(foo2_func)

        assert not (foo2_func in model.statements)
        assert foo1_func in model.statements