when doing line-matching, check for # This is a test template in the line

Author: joyguoguo

fuzz/collect_fuzz_python.py

@@ -303,23 +303,33 @@ def visit_FunctionDef(self, node):
         return node
 
     def add_param_assignment(self, node, param_name):
-        """Add param_name = b"" at the beginning of the function body"""
+        """Add param_name = b"..." at the beginning of the function body with an inline comment"""
+        # 创建包含赋值和注释的复合值
+        value_with_comment = ast.JoinedStr(
+            values=[
+                ast.FormattedValue(value=ast.Constant(value=b""), conversion=-1),
+                ast.Constant(value="  # This is a test template")
+            ]
+        )
+        
         # 创建赋值节点
         assign_node = ast.Assign(
             targets=[ast.Name(id=param_name, ctx=ast.Store())],
-            value=ast.Constant(value=b""),
+            value=value_with_comment
         )
-
+        
         # 如果有文档字符串，插入在文档字符串之后
         if (
             node.body
             and isinstance(node.body[0], ast.Expr)
-            and isinstance(node.body[0].value, ast.Str)
+            and isinstance(node.body[0].value, ast.Constant)
+            and isinstance(node.body[0].value.value, str)
         ):
+            # 插入在文档字符串后面
             node.body.insert(1, assign_node)
         else:
+            # 插入在函数开头
             node.body.insert(0, assign_node)
-
     def remove_print_param(self, node, param_name):
         """Remove print statements for the specific parameter"""
         new_body = []
@@ -360,62 +370,30 @@ def visit_If(self, node):
 
 
 class TestGenTransformer(ast.NodeTransformer):
-    """AST transformer for generating test cases from fuzzing inputs"""
-
-    def __init__(self, idx: int, fuzz_input: bytes):
+    def __init__(self, idx, fuzz_input):
         self.idx = idx
         self.fuzz_input = fuzz_input
         self.found_test_function = False
 
     def visit_FunctionDef(self, node):
-        # 只处理名为 test_ 的函数
         if node.name == "test_":
             self.found_test_function = True
-
-            # 1. 将函数名改为 test_{idx}
-            node.name = f"test_{self.idx}"
-
-            # 2. 找到并替换 data = b"" 赋值语句
-            self.replace_data_assignment(node)
-
+            # 遍历函数体，寻找包含注释的赋值语句
+            for i, stmt in enumerate(node.body):
+                # 检查是否是赋值语句
+                if isinstance(stmt, ast.Assign):
+                    # 检查赋值语句的值是否是带有注释的复合值
+                    if (
+                        isinstance(stmt.value, ast.JoinedStr)
+                        and len(stmt.value.values) >= 2
+                        and isinstance(stmt.value.values[1], ast.Constant)
+                        and stmt.value.values[1].value == "  # This is a test template"
+                    ):
+                        # 替换为新的输入值
+                        stmt.value = ast.Constant(value=self.fuzz_input)
+                        break
         return node
 
-    def replace_data_assignment(self, node):
-        """Replace data assignment with fuzz input"""
-        for i, stmt in enumerate(node.body):
-            # 查找赋值语句
-            if isinstance(stmt, ast.Assign):
-                # 检查是否是 data = b"" 格式的赋值
-                if (
-                    len(stmt.targets) == 1
-                    and isinstance(stmt.targets[0], ast.Name)
-                    and isinstance(stmt.value, ast.Constant)
-                    and stmt.value.value == b""
-                ):
-
-                    # 替换为新的输入数据
-                    node.body[i] = ast.Assign(
-                        targets=[stmt.targets[0]],
-                        value=ast.Constant(value=self.fuzz_input),
-                    )
-                    return
-
-                # 检查是否是 data = b'' 格式的赋值
-                if (
-                    len(stmt.targets) == 1
-                    and isinstance(stmt.targets[0], ast.Name)
-                    and isinstance(stmt.value, ast.Constant)
-                    and stmt.value.value == b""
-                ):
-
-                    # 替换为新的输入数据
-                    node.body[i] = ast.Assign(
-                        targets=[stmt.targets[0]],
-                        value=ast.Constant(value=self.fuzz_input),
-                    )
-                    return
-
-
 def substitute_one_repo(
     repo: str,
     targets: list[str],