better test, improve testing

Author: aseembits93

codeflash/code_utils/lprof_utils.py

@@ -23,38 +23,38 @@ def leave_FunctionDef(self, original_node, updated_node):
     updated_module = module_node.visit(transformer)
     return updated_module
 
-def add_decorator_imports(file_paths, fn_list, db_file):
-    """Adds a decorator to a function in a Python file."""
-    for file_path, fn_name in zip(file_paths, fn_list):
-        #open file
-        with open(file_path, "r", encoding="utf-8") as file:
-            file_contents = file.read()
+def add_profile_enable(original_code: str, db_file: str) -> str:
+    module = cst.parse_module(original_code)
+    found_index = -1
+
+    for idx, statement in enumerate(module.body):
+        if isinstance(statement, cst.SimpleStatementLine):
+            for stmt in statement.body:
+                if isinstance(stmt, cst.ImportFrom):
+                    if stmt.module and stmt.module.value == 'line_profiler':
+                        for name in stmt.names:
+                            if isinstance(name, cst.ImportAlias):
+                                if name.name.value == 'profile' and name.asname is None:
+                                    found_index = idx
+                                    break
+                        if found_index != -1:
+                            break
+        if found_index != -1:
+            break
 
-        # parse to cst
-        module_node = cst.parse_module(file_contents)
-        # add decorator
-        module_node = add_decorator_cst(module_node, fn_name, 'profile')
-        # add imports
-        # Create a transformer to add the import
-        transformer = ImportAdder("from line_profiler import profile")
+    if found_index == -1:
+        return original_code  # or raise an exception if the import is not found
 
-        # Apply the transformer to add the import
-        module_node = module_node.visit(transformer)
-        modified_code = isort.code(module_node.code, float_to_top=True)
-        # write to file
-        with open(file_path, "w", encoding="utf-8") as file:
-            file.write(modified_code)
-    #Adding profile.enable line for changing the savepath of the data, do this only for the main file and not the helper files, can use libcst seems like an overkill, will go just with some simple string manipulation
-    with open(file_paths[0],'r') as f:
-        file_contents = f.readlines()
-    for idx, line in enumerate(file_contents):
-        if 'from line_profiler import profile' in line:
-            file_contents.insert(idx+1, f"profile.enable(output_prefix='{db_file}')\n")
-            break
-    with open(file_paths[0],'w') as f:
-        f.writelines(file_contents)
+    # Create the new line to insert
+    new_line = f"profile.enable(output_prefix='{db_file}')\n"
+    new_statement = cst.parse_statement(new_line)
 
+    # Insert the new statement into the module's body
+    new_body = list(module.body)
+    new_body.insert(found_index + 1, new_statement)
+    modified_module = module.with_changes(body=new_body)
 
+    return modified_module.code
 
 
 class ImportAdder(cst.CSTTransformer):
@@ -82,6 +82,36 @@ def visit_ImportFrom(self, node):
                 if import_alias.name.value == "profile":
                     self.has_import = True
 
+
+def add_decorator_imports(file_paths, fn_list, db_file):
+    """Adds a decorator to a function in a Python file."""
+    for file_path, fn_name in zip(file_paths, fn_list):
+        #open file
+        with open(file_path, "r", encoding="utf-8") as file:
+            file_contents = file.read()
+
+        # parse to cst
+        module_node = cst.parse_module(file_contents)
+        # add decorator
+        module_node = add_decorator_cst(module_node, fn_name, 'profile')
+        # add imports
+        # Create a transformer to add the import
+        transformer = ImportAdder("from line_profiler import profile")
+
+        # Apply the transformer to add the import
+        module_node = module_node.visit(transformer)
+        modified_code = isort.code(module_node.code, float_to_top=True)
+        # write to file
+        with open(file_path, "w", encoding="utf-8") as file:
+            file.write(modified_code)
+    #Adding profile.enable line for changing the savepath of the data, do this only for the main file and not the helper files, can use libcst seems like an overkill, will go just with some simple string manipulation
+    with open(file_paths[0],'r') as f:
+        file_contents = f.read()
+    modified_code = add_profile_enable(file_contents,db_file)
+    with open(file_paths[0],'w') as f:
+        f.write(modified_code)
+
+
 def prepare_lprofiler_files(prefix: str = "") -> tuple[Path]:
     """Prepare line profiler output file."""
     lprofiler_database_file = get_run_tmp_file(Path(prefix))

tests/test_test_runner.py

@@ -8,6 +8,7 @@
 from codeflash.verification.test_results import TestType
 from codeflash.verification.test_runner import run_behavioral_tests
 from codeflash.verification.verification_utils import TestConfig
+import dill as pickle
 
 
 def test_unittest_runner():
@@ -144,11 +145,18 @@ def test_sort():
             pytest_target_runtime_seconds=1,
             enable_lprofiler=True,
         )
-        with open(tmpdir.name+os.sep+"/baseline.txt", "r") as f:
-            output = f.read()
-        expected_output = "Timer unit: 1e-09 s\n\nTotal time: 0 s\nFile: {}\nFunction: sorter at line 4\n\nLine #      Hits         Time  Per Hit   % Time  Line Contents\n==============================================================\n     4                                           @profile\n     5                                           def sorter(arr):\n     6         1          0.0      0.0               arr.sort()\n     7         1          0.0      0.0               return arr\n\n".format(fp.name)
-
-
-    assert output == expected_output, "Test passed"
+        try:
+            #todo write a test for the pickle parsing code, right now it's a simplistic test
+            with open(tmpdir.name+os.sep+"baseline.lprof", "rb") as f:
+                output = pickle.load(f)
+            #get lprof instead and compare the hits
+            output_set = set(list(output.timings.values())[0])
+            output_set = {(x,y) for x,y,z in output_set}
+            expected_set = {(6, 1), (7, 1)}
+        except Exception as e:
+            print(e)
+            assert False, "Test failed"
+
+    assert expected_set == output_set, "Test passed"
     result_file.unlink(missing_ok=True)