new tests

Author: aseembits93

codeflash/code_utils/code_replacer.py

@@ -7,7 +7,6 @@
 
 import isort
 import libcst as cst
-import libcst.matchers as m
 from libcst.metadata import PositionProvider
 
 from codeflash.cli_cmds.console import logger
@@ -41,29 +40,50 @@ def normalize_code(code: str) -> str:
 class AddRequestArgument(cst.CSTTransformer):
     METADATA_DEPENDENCIES = (PositionProvider,)
 
-    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:  # noqa: ARG002
-        args = updated_node.params.params
-        arg_names = {arg.name.value for arg in args}
-
-        # Skip if 'request' is already present
-        if "request" in arg_names:
-            return updated_node
-
-        # Create a new 'request' param
-        request_param = cst.Param(name=cst.Name("request"))
-
-        # Add 'request' as the first argument (after 'self' or 'cls' if needed)
-        if args:
-            first_arg = args[0].name.value
-            if first_arg in {"self", "cls"}:
-                new_params = [args[0], request_param] + list(args[1:])  # noqa: RUF005
-            else:
-                new_params = [request_param] + list(args)  # noqa: RUF005
-        else:
-            new_params = [request_param]
-
-        new_param_list = updated_node.params.with_changes(params=new_params)
-        return updated_node.with_changes(params=new_param_list)
+    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
+        # Matcher for '@fixture' or '@pytest.fixture'
+        for decorator in original_node.decorators:
+            dec = decorator.decorator
+
+            if isinstance(dec, cst.Call):
+                func_name = ""
+                if isinstance(dec.func, cst.Attribute) and isinstance(dec.func.value, cst.Name):
+                    if dec.func.attr.value == "fixture" and dec.func.value.value == "pytest":
+                        func_name = "pytest.fixture"
+                elif isinstance(dec.func, cst.Name) and dec.func.value == "fixture":
+                    func_name = "fixture"
+
+                if func_name:
+                    for arg in dec.args:
+                        if (
+                            arg.keyword
+                            and arg.keyword.value == "autouse"
+                            and isinstance(arg.value, cst.Name)
+                            and arg.value.value == "True"
+                        ):
+                            args = updated_node.params.params
+                            arg_names = {arg.name.value for arg in args}
+
+                            # Skip if 'request' is already present
+                            if "request" in arg_names:
+                                return updated_node
+
+                            # Create a new 'request' param
+                            request_param = cst.Param(name=cst.Name("request"))
+
+                            # Add 'request' as the first argument (after 'self' or 'cls' if needed)
+                            if args:
+                                first_arg = args[0].name.value
+                                if first_arg in {"self", "cls"}:
+                                    new_params = [args[0], request_param] + list(args[1:])  # noqa: RUF005
+                                else:
+                                    new_params = [request_param] + list(args)  # noqa: RUF005
+                            else:
+                                new_params = [request_param]
+
+                            new_param_list = updated_node.params.with_changes(params=new_params)
+                            return updated_node.with_changes(params=new_param_list)
+        return updated_node
 
 
 class PytestMarkAdder(cst.CSTTransformer):
@@ -135,33 +155,41 @@ def _create_pytest_mark(self) -> cst.Decorator:
 class AutouseFixtureModifier(cst.CSTTransformer):
     def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
         # Matcher for '@fixture' or '@pytest.fixture'
-        fixture_decorator_func = m.Name("fixture") | m.Attribute(value=m.Name("pytest"), attr=m.Name("fixture"))
-
         for decorator in original_node.decorators:
-            if m.matches(
-                decorator,
-                m.Decorator(
-                    decorator=m.Call(
-                        func=fixture_decorator_func, args=[m.Arg(value=m.Name("True"), keyword=m.Name("autouse"))]
-                    )
-                ),
-            ):
-                # Found a matching fixture with autouse=True
-
-                # 1. The original body of the function will become the 'else' block.
-                #    updated_node.body is an IndentedBlock, which is what cst.Else expects.
-                else_block = cst.Else(body=updated_node.body)
-
-                # 2. Create the new 'if' block that will exit the fixture early.
-                if_test = cst.parse_expression('request.node.get_closest_marker("codeflash_no_autouse")')
-                yield_statement = cst.parse_statement("yield")
-                if_body = cst.IndentedBlock(body=[yield_statement])
-
-                # 3. Construct the full if/else statement.
-                new_if_statement = cst.If(test=if_test, body=if_body, orelse=else_block)
-
-                # 4. Replace the entire function's body with our new single statement.
-                return updated_node.with_changes(body=cst.IndentedBlock(body=[new_if_statement]))
+            dec = decorator.decorator
+
+            if isinstance(dec, cst.Call):
+                func_name = ""
+                if isinstance(dec.func, cst.Attribute) and isinstance(dec.func.value, cst.Name):
+                    if dec.func.attr.value == "fixture" and dec.func.value.value == "pytest":
+                        func_name = "pytest.fixture"
+                elif isinstance(dec.func, cst.Name) and dec.func.value == "fixture":
+                    func_name = "fixture"
+
+                if func_name:
+                    for arg in dec.args:
+                        if (
+                            arg.keyword
+                            and arg.keyword.value == "autouse"
+                            and isinstance(arg.value, cst.Name)
+                            and arg.value.value == "True"
+                        ):
+                            # Found a matching fixture with autouse=True
+
+                            # 1. The original body of the function will become the 'else' block.
+                            #    updated_node.body is an IndentedBlock, which is what cst.Else expects.
+                            else_block = cst.Else(body=updated_node.body)
+
+                            # 2. Create the new 'if' block that will exit the fixture early.
+                            if_test = cst.parse_expression('request.node.get_closest_marker("codeflash_no_autouse")')
+                            yield_statement = cst.parse_statement("yield")
+                            if_body = cst.IndentedBlock(body=[yield_statement])
+
+                            # 3. Construct the full if/else statement.
+                            new_if_statement = cst.If(test=if_test, body=if_body, orelse=else_block)
+
+                            # 4. Replace the entire function's body with our new single statement.
+                            return updated_node.with_changes(body=cst.IndentedBlock(body=[new_if_statement]))
         return updated_node