Refactor the module creation and clean up

Author: vijaykanthm

securitycenter/snippets_management_api/security_health_analytics_custom_modules.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 #
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

securitycenter/snippets_management_api/security_health_analytics_custom_modules_test.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 #
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,6 +19,8 @@
 
 import time
 
+import uuid
+
 import backoff
 
 from google.api_core.exceptions import InternalServerError, NotFound, ServiceUnavailable
@@ -33,30 +35,103 @@
 # GCLOUD_ORGANIZATION: The organization ID.
 ORGANIZATION_ID = os.environ["GCLOUD_ORGANIZATION"]
 LOCATION = "global"
-PREFIX = "python_sample_sha_custom_module"  # Prefix used for identifying test modules
+PREFIX = "python_sample_sha_custom_module"
 
-# Global list to track created modules
-created_modules = []
+# Global list to track created shared modules
+shared_modules = []
 
 
 @pytest.fixture(scope="session", autouse=True)
 def setup_environment():
     if not ORGANIZATION_ID:
         pytest.fail("GCLOUD_ORGANIZATION environment variable is not set.")
 
+    setup_shared_modules()
+
 
 @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()
+        print_all_shared_modules()
+        cleanup_shared_modules()
 
     request.addfinalizer(teardown)
 
 
+def setup_shared_modules():
+    for _ in range(3) :
+        _, module_id = add_custom_module(ORGANIZATION_ID)
+        if module_id != "" :
+            shared_modules.append(module_id)
+
+
+def add_module_to_cleanup(module_id):
+    shared_modules.append(module_id)
+
+
+def print_all_shared_modules():
+    """Print all created custom modules."""
+    if not shared_modules:
+        print("No custom modules were created.")
+    else:
+        print("\nCreated Custom Modules:")
+        for module_id in shared_modules:
+            print(module_id)
+
+
+def cleanup_shared_modules():
+    """
+    Deletes all created custom modules in this test session.
+    """
+    client = securitycentermanagement_v1.SecurityCenterManagementClient()
+
+    print("Cleaning up created custom modules...")
+
+    for module_id in list(shared_modules):
+        if not custom_module_exists(module_id):
+            print(f"Module not found (already deleted): {module_id}")
+            shared_modules.remove(module_id)
+            continue
+        try:
+            client.delete_security_health_analytics_custom_module(
+                    request={"name": f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}/securityHealthAnalyticsCustomModules/{module_id}"}
+                )
+            print(f"Deleted custom module: {module_id}")
+            shared_modules.remove(module_id)
+        except Exception as e:
+            print(f"Failed to delete module {module_id}: {e}")
+            raise
+
+
+def custom_module_exists(module_id):
+    client = securitycentermanagement_v1.SecurityCenterManagementClient()
+    try:
+        client.get_security_health_analytics_custom_module(
+                request={"name": f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}/securityHealthAnalyticsCustomModules/{module_id}"}
+            )
+        return True
+    except Exception as e:
+        if "404" in str(e):
+            return False
+        raise
+
+
+def get_random_shared_module():
+    if not shared_modules:
+        return ""
+    random.seed(int(time.time() * 1000000))
+    return shared_modules[random.randint(0, len(shared_modules) - 1)]
+
+
+def extract_custom_module_id(module_name):
+    trimmed_full_name = module_name.strip()
+    parts = trimmed_full_name.split("/")
+    if parts:
+        return parts[-1]
+    return ""
+
+
 def add_custom_module(org_id: str):
     """
     Adds a new SHA custom module.
@@ -69,7 +144,7 @@ def add_custom_module(org_id: str):
     client = securitycentermanagement_v1.SecurityCenterManagementClient()
 
     # Generate a unique display name
-    unique_suffix = f"{int(time.time())}_{random.randint(0, 999)}"
+    unique_suffix = str(uuid.uuid4()).replace("-", "_")
     display_name = f"python_sample_sha_custom_module_test_{unique_suffix}"
 
     # Define the custom module configuration
@@ -109,7 +184,7 @@ def add_custom_module(org_id: str):
     response = client.create_security_health_analytics_custom_module(request=request)
     print(f"Created Security Health Analytics Custom Module: {response.name}")
     module_name = response.name
-    module_id = module_name.split("/")[-1]
+    module_id = extract_custom_module_id(module_name)
     return module_name, module_id
 
 
@@ -118,8 +193,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)
+    module_id = get_random_shared_module()
     parent = f"organizations/{ORGANIZATION_ID}/locations/{LOCATION}"
 
     # Retrieve the custom module
@@ -129,68 +203,39 @@ def test_get_effective_security_health_analytics_custom_module():
     # Verify that the custom module was created
     assert response.display_name.startswith(PREFIX)
     assert response.enablement_state == securitycentermanagement_v1.EffectiveSecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
-    print(f"Retrieved Custom Module: {response.name}")
 
 
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
 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)
 
     assert custom_modules is not None, "Failed to retrieve the custom modules."
     assert len(custom_modules) > 0, "No custom modules were retrieved."
 
-    # Verify the created module is in the list
-    created_module = next(
-        (module for module in custom_modules if module.name == module_name), None
-    )
-    assert created_module is not None, "Created custom module not found in the list."
-    assert created_module.display_name.startswith(PREFIX)
-    assert (
-        created_module.enablement_state
-        == securitycentermanagement_v1.SecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
-    )
-
 
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
 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)
 
     assert custom_modules is not None, "Failed to retrieve the custom modules."
     assert len(custom_modules) > 0, "No custom modules were retrieved."
 
-    # Verify the created module is in the list
-    created_module = next(
-        (module for module in custom_modules if (module.name.split("/")[-1]) == module_id), None
-    )
-    assert created_module is not None, "Created custom module not found in the list."
-    assert created_module.display_name.startswith(PREFIX)
-    assert (
-        created_module.enablement_state
-        == securitycentermanagement_v1.EffectiveSecurityHealthAnalyticsCustomModule.EnablementState.ENABLED
-    )
-
 
 @backoff.on_exception(
     backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
 )
 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)
@@ -199,47 +244,3 @@ 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