diff --git a/tests/unit/ast/test_annotate_and_optimize_ast.py b/tests/unit/ast/test_annotate_and_optimize_ast.py
index 936c81e21a..39ea899bd9 100644
--- a/tests/unit/ast/test_annotate_and_optimize_ast.py
+++ b/tests/unit/ast/test_annotate_and_optimize_ast.py
@@ -31,7 +31,7 @@ def get_contract_info(source_code):
     pre_parse_result = pre_parse(source_code)
     py_ast = python_ast.parse(pre_parse_result.reformatted_code)
 
-    annotate_python_ast(py_ast, pre_parse_result)
+    annotate_python_ast(py_ast, pre_parse_result.reformatted_code, pre_parse_result)
 
     return py_ast, pre_parse_result.reformatted_code