1 files changed, 15 insertions, 6 deletions
diff --git a/tests/atf_python/atf_pytest.py b/tests/atf_python/atf_pytest.py
index db7244d3234b..0dd3a225b73d 100644
--- a/tests/atf_python/atf_pytest.py
+++ b/tests/atf_python/atf_pytest.py
@@ -6,15 +6,12 @@ from typing import NamedTuple
 from typing import Optional
 from typing import Tuple
 
+from atf_python.utils import nodeid_to_method_name
+
 import pytest
 import os
 
 
-def nodeid_to_method_name(nodeid: str) -> str:
-    """file_name.py::ClassName::method_name[parametrize] -> method_name"""
-    return nodeid.split("::")[-1].split("[")[0]
-
-
 class ATFCleanupItem(pytest.Item):
     def runtest(self):
         """Runs cleanup procedure for the test instead of the test itself"""
@@ -73,7 +70,6 @@ class ATFTestObj(object):
         else:
             ret["require.user"] = username
 
-
     def _convert_marks(self, obj) -> Dict[str, Any]:
         wj_func = lambda x: " ".join(x)  # noqa: E731
         _map: Dict[str, Dict] = {
@@ -158,6 +154,19 @@ class ATFHandler(object):
                 return True
         return False
 
+    def _generate_test_cleanups(self, items):
+        new_items = []
+        for obj in items:
+            if self.has_object_cleanup(obj):
+                self.override_runtest(obj)
+                new_items.append(obj)
+        items.clear()
+        items.extend(new_items)
+
+    def modify_tests(self, items, config):
+        if config.option.atf_cleanup:
+            self._generate_test_cleanups(items)
+
     def list_tests(self, tests: List[str]):
         print('Content-Type: application/X-atf-tp; version="1"')
         print()