Clarified docstring for get_modification_code_ranges

... and started adding tests for that function.

Author: zomglings

codeflash/code_utils/formatter.py

@@ -60,15 +60,14 @@ def sort_imports(code: str) -> str:
 
     return sorted_code
 
-# TODO(zomglings): Write unit tests.
 def get_modification_code_ranges(
     modified_code: str,
     fto: FunctionToOptimize,
     preexisting_functions: set[tuple[str, tuple[FunctionParent,...]]],
     helper_functions: list[FunctionSource],
 ) -> list[tuple[int, int]]:
     """
-    Returns the line number of modified and new functions in a string containing containing the code in a fully modified file.
+    Returns the starting and ending line numbers of modified and new functions in a file containing edits.
     """
     modified_functions = set()
     modified_functions.add(fto.qualified_name)

tests/test_formatter.py

@@ -5,7 +5,8 @@
 import pytest
 
 from codeflash.code_utils.config_parser import parse_config_file
-from codeflash.code_utils.formatter import format_code, sort_imports
+from codeflash.code_utils.formatter import format_code, get_modification_code_ranges, sort_imports
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 
 
 def test_remove_duplicate_imports():
@@ -209,3 +210,15 @@ def foo():
         tmp_path = tmp.name
         with pytest.raises(FileNotFoundError):
             format_code(formatter_cmds=["exit 1"], path=Path(tmp_path))
+
+def test_get_modification_code_ranges_self_contained_fto():
+    modified_code = """
+def hello(name):
+    print(f"Hello, {{name}}")
+"""
+
+    fto = FunctionToOptimize(function_name="hello", file_path=Path("hello.py"), parents=[])
+    code_ranges = get_modification_code_ranges(modified_code, fto, set(), [])
+
+    assert len(code_ranges) == 1
+    assert code_ranges[0] == (2, 3)