fix import pytest later

Tesorio · Sep 17, 2024 · f05becd · f05becd
1 parent 39571d0
commit f05becd
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/pytest_xflaky/add_decorator.py b/src/pytest_xflaky/add_decorator.py
@@ -59,11 +59,6 @@ def traverse(node):
     else:
         return None
 
-    # add import pytest
-    if not is_pytest_imported:
-        import_statement = "import pytest\n"
-        source_code = import_statement + source_code
-
     # skip if decorator already added
     if not any(d.startswith(b"@pytest.mark.xfail") for d in decorators):
         # Add the decorator before the function definition
@@ -75,6 +70,11 @@ def traverse(node):
             + source_code[function_start_byte:]
         )
 
+    # add import pytest
+    if not is_pytest_imported:
+        import_statement = "import pytest\n"
+        source_code = import_statement + source_code
+
     with open(path, "w") as fp:
         fp.write(source_code)