reworked tests

Author: alvin-r

codeflash/context/code_context_extractor.py

@@ -183,14 +183,16 @@ def extract_code_string_context_from_files(
             helpers_of_helpers_qualified_names = {
                 func.qualified_name for func in helpers_of_helpers.get(file_path, set())
             }
+            code_without_unused_defs = remove_unused_definitions_by_function_names(
+                original_code, qualified_function_names | helpers_of_helpers_qualified_names
+            )
             code_context = parse_code_and_prune_cst(
-                original_code,
+                code_without_unused_defs,
                 code_context_type,
                 qualified_function_names,
                 helpers_of_helpers_qualified_names,
                 remove_docstrings,
             )
-            code_context = remove_unused_definitions_by_function_names(code_context, qualified_function_names | helpers_of_helpers_qualified_names)
         except ValueError as e:
             logger.debug(f"Error while getting read-only code: {e}")
             continue
@@ -215,10 +217,10 @@ def extract_code_string_context_from_files(
             continue
         try:
             qualified_helper_function_names = {func.qualified_name for func in helper_function_sources}
+            code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, qualified_helper_function_names)
             code_context = parse_code_and_prune_cst(
-                original_code, code_context_type, set(), qualified_helper_function_names, remove_docstrings
+                code_without_unused_defs, code_context_type, set(), qualified_helper_function_names, remove_docstrings
             )
-            code_context = remove_unused_definitions_by_function_names(code_context, qualified_helper_function_names)
         except ValueError as e:
             logger.debug(f"Error while getting read-only code: {e}")
             continue
@@ -285,16 +287,16 @@ def extract_code_markdown_context_from_files(
             helpers_of_helpers_qualified_names = {
                 func.qualified_name for func in helpers_of_helpers.get(file_path, set())
             }
+            code_without_unused_defs = remove_unused_definitions_by_function_names(
+                original_code, qualified_function_names | helpers_of_helpers_qualified_names
+            )
             code_context = parse_code_and_prune_cst(
-                original_code,
+                code_without_unused_defs,
                 code_context_type,
                 qualified_function_names,
                 helpers_of_helpers_qualified_names,
                 remove_docstrings,
             )
-            code_context = remove_unused_definitions_by_function_names(
-                code_context, qualified_function_names | helpers_of_helpers_qualified_names
-            )
 
         except ValueError as e:
             logger.debug(f"Error while getting read-only code: {e}")
@@ -323,11 +325,9 @@ def extract_code_markdown_context_from_files(
             continue
         try:
             qualified_helper_function_names = {func.qualified_name for func in helper_function_sources}
+            code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, qualified_helper_function_names)
             code_context = parse_code_and_prune_cst(
-                original_code, code_context_type, set(), qualified_helper_function_names, remove_docstrings
-            )
-            code_context = remove_unused_definitions_by_function_names(
-                code_context, qualified_helper_function_names
+                code_without_unused_defs, code_context_type, set(), qualified_helper_function_names, remove_docstrings
             )
         except ValueError as e:
             logger.debug(f"Error while getting read-only code: {e}")