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
