Refactor the cleaning up of created custom modules in the current test session

Author: vijaykanthm

securitycenter/snippets_management_api/security_health_analytics_custom_modules_test.py

@@ -35,46 +35,26 @@
 LOCATION = "global"
 PREFIX = "python_sample_sha_custom_module"  # Prefix used for identifying test modules
 
+# Global list to track created modules
+created_modules = []
+
 
 @pytest.fixture(scope="session", autouse=True)
 def setup_environment():
-    """
-    Fixture to ensure a clean environment by removing test modules before running tests.
-
-    This fixture lists all SHA custom modules in the organization and deletes any
-    that were created by previous test runs, identified by the PREFIX.
-    """
     if not ORGANIZATION_ID:
         pytest.fail("GCLOUD_ORGANIZATION environment variable is not set.")
 
-    print(f"Cleaning up existing custom modules for organization: {ORGANIZATION_ID}")
-    cleanup_existing_custom_modules(ORGANIZATION_ID)
 
+@pytest.fixture(scope="session", autouse=True)
+def cleanup_after_tests(request):
+    """Fixture to clean up created custom modules after the test session."""
+    def teardown():
+        print("\nCreated Custom Modules:")
+        print_all_created_modules()
+        print("Cleaning up created custom modules...")
+        cleanup_created_custom_modules()
 
-def cleanup_existing_custom_modules(org_id: str):
-    """
-    Deletes all custom modules matching a specific naming pattern.
-    Args:
-        org_id: The organization ID.
-    """
-    client = securitycentermanagement_v1.SecurityCenterManagementClient()
-    parent = f"organizations/{org_id}/locations/global"
-    print(f"Parent path: {parent}")
-    try:
-        custom_modules = client.list_security_health_analytics_custom_modules(
-            request={"parent": parent}
-        )
-        for module in custom_modules:
-            if module.display_name.startswith(PREFIX):
-                client.delete_security_health_analytics_custom_module(
-                    request={"name": module.name}
-                )
-                print(f"Deleted custom module: {module.name}")
-    except NotFound as e:
-        print(f"Resource not found: {e}")
-    except Exception as e:
-        print(f"Unexpected error during cleanup: {e}")
-        raise
+    request.addfinalizer(teardown)
 
 
 def add_custom_module(org_id: str):
@@ -139,6 +119,7 @@ def add_custom_module(org_id: str):
 def test_get_effective_security_health_analytics_custom_module():
     """Tests getting an effective SHA custom module."""
     module_name, module_id = add_custom_module(ORGANIZATION_ID)
+    created_modules.append(module_name)
     parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
 
     # Retrieve the custom module
@@ -157,6 +138,7 @@ def test_get_effective_security_health_analytics_custom_module():
 def test_list_descendant_security_health_analytics_custom_module():
     """Tests listing descendant SHA custom modules."""
     module_name, module_id = add_custom_module(ORGANIZATION_ID)
+    created_modules.append(module_name)
     parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
     # Retrieve the list descendant custom modules
     custom_modules = security_health_analytics_custom_modules.list_descendant_security_health_analytics_custom_module(parent)
@@ -182,6 +164,7 @@ def test_list_descendant_security_health_analytics_custom_module():
 def test_list_effective_security_health_analytics_custom_module():
     """Tests listing effective SHA custom modules."""
     module_name, module_id = add_custom_module(ORGANIZATION_ID)
+    created_modules.append(module_name)
     parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
     # Retrieve the list of custom modules
     custom_modules = security_health_analytics_custom_modules.list_effective_security_health_analytics_custom_module(parent)
@@ -207,6 +190,7 @@ def test_list_effective_security_health_analytics_custom_module():
 def test_simulate_security_health_analytics_custom_module():
     """Tests simulating an SHA custom module."""
     module_name, module_id = add_custom_module(ORGANIZATION_ID)
+    created_modules.append(module_name)
     parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
 
     simulated_custom_module = security_health_analytics_custom_modules.simulate_security_health_analytics_custom_module(parent)
@@ -215,3 +199,47 @@ def test_simulate_security_health_analytics_custom_module():
     assert simulated_custom_module.result.no_violation is not None, (
         f"Expected no_violation to be present, got {simulated_custom_module.result}."
     )
+
+
+def print_all_created_modules():
+    """Print all created custom modules."""
+    if not created_modules:
+        print("No custom modules were created.")
+    else:
+        for module in created_modules:
+            print(module)
+
+
+def cleanup_created_custom_modules():
+    """
+    Deletes all created custom modules in this test session.
+    """
+    client = securitycentermanagement_v1.SecurityCenterManagementClient()
+
+    for module in list(created_modules):
+        if not custom_module_exists(module):
+            print(f"Module not found (already deleted): {module}")
+            created_modules.remove(module)
+            continue
+        try:
+            client.delete_security_health_analytics_custom_module(
+                    request={"name": module}
+                )
+            print(f"Deleted custom module: {module}")
+            created_modules.remove(module)
+        except Exception as e:
+            print(f"Failed to delete module {module}: {e}")
+            raise
+
+
+def custom_module_exists(module_name):
+    client = securitycentermanagement_v1.SecurityCenterManagementClient()
+    try:
+        client.get_security_health_analytics_custom_module(
+                request={"name": module_name}
+            )
+        return True
+    except Exception as e:
+        if "404" in str(e):
+            return False
+        raise