SQUASH Quarantine logic

Author: gmarciani

tests/integration-tests/constants.py

@@ -19,6 +19,15 @@
 
 UNSUPPORTED_OSES_FOR_DCV = ["alinux2023"]
 
+DO_NOT_DELETE_TAG_KEY = 'DO-NOT-DELETE'
+
+QUARANTINE_TAG_KEY = 'QUARANTINE'
+
+MAX_QUARANTINED_STACKS = 5
+
+QUARANTINE_TAGS = [{ "Key": DO_NOT_DELETE_TAG_KEY, "Value": "true" }, { "Key": QUARANTINE_TAG_KEY, "Value": "true" }]
+RETENTION_TAGS = [{ "Key": DO_NOT_DELETE_TAG_KEY, "Value": "true" }]
+
 
 class NodeType(Enum):
     """Categories of nodes."""

tests/integration-tests/tests/ad_integration/test_ad_integration.py

@@ -28,20 +28,17 @@
 from retrying import retry
 from time_utils import seconds
 
-from utils import find_stack_by_tag, generate_stack_name, is_directory_supported, random_alphanumeric
+from utils import find_stack_by_tag, generate_stack_name, is_directory_supported, random_alphanumeric, get_quarantined_stacks, is_quarantined_stack, quarantine_stacks
 
 from tests.ad_integration.cluster_user import ClusterUser
 from tests.common.utils import run_system_analyzer
+from constants import DO_NOT_DELETE_TAG_KEY, MAX_QUARANTINED_STACKS
 
 NUM_USERS_TO_CREATE = 5
 NUM_USERS_TO_TEST = 3
 
-MAX_QUARANTINED_STACKS = 5
-
 AD_STACK_PREFIX = 'integ-tests-MultiUserInfraStack'
 
-DO_NOT_DELETE_TAG_KEY = 'DO-NOT-DELETE'
-
 
 def get_infra_stack_outputs(stack_name):
     cfn = boto3.client("cloudformation")
@@ -228,15 +225,14 @@ def _create_directory_stack(cfn_stacks_factory, request, directory_type, region,
         logging.error("Failed to create stack %s", directory_stack_name)
         # We want to retain the stack in case of failure in order to debug it.
         # We retain a limited number of stack to contain the costs.
-        n_retained_ad_stacks = get_retained_ad_stacks_count()
-        if n_retained_ad_stacks < MAX_QUARANTINED_STACKS:
-            logging.warn("Retaining failed stack %s to debug failure (retained: %d, max: %d)",
-                         directory_stack_name, n_retained_ad_stacks, MAX_QUARANTINED_STACKS)
-
-            add_tag_to_stack(directory_stack.name, DO_NOT_DELETE_TAG_KEY, "Retained to debug failure")
+        n_quarantined_ad_stacks = len(get_quarantined_stacks(region, prefix=AD_STACK_PREFIX))
+        if n_quarantined_ad_stacks < MAX_QUARANTINED_STACKS:
+            logging.warn("Quarantining failed stack %s to debug failure (quarantined: %d, max: %d)",
+                         directory_stack_name, n_quarantined_ad_stacks, MAX_QUARANTINED_STACKS)
+            quarantine_stacks(region, stack_names=[directory_stack_name])
         else:
-            logging.warn("Cannot retain failed stack %s for debugging because there are already %d retained (max: %d)",
-                         directory_stack_name, n_retained_ad_stacks, MAX_QUARANTINED_STACKS)
+            logging.warn("Cannot quarantine failed stack %s for debugging because there are already %d quarantined (max: %d)",
+                         directory_stack_name, n_quarantined_ad_stacks, MAX_QUARANTINED_STACKS)
         raise e
     logging.info("Creation of stack %s complete", directory_stack_name)
 
@@ -276,7 +272,7 @@ def _directory_factory(
             stack_prefix = f"{AD_STACK_PREFIX}{directory_type}"
             directory_stack_name = find_stack_by_tag("parallelcluster:integ-tests-ad-stack", region, stack_prefix)
 
-            if not directory_stack_name:
+            if not directory_stack_name or is_quarantined_stack(region, directory_stack_name):
                 directory_stack = _create_directory_stack(
                     cfn_stacks_factory,
                     request,

tests/integration-tests/utils.py

@@ -28,6 +28,7 @@
 from jinja2.sandbox import SandboxedEnvironment
 from retrying import retry
 from time_utils import minutes, seconds
+from constants import QUARANTINE_TAG_KEY, DO_NOT_DELETE_TAG_KEY, QUARANTINE_TAGS
 
 DEFAULT_PARTITION = "aws"
 PARTITION_MAP = {
@@ -908,3 +909,44 @@ def find_stack_by_tag(tag, region, stack_prefix):
             logging.info(f"Found stack: {name} (created on {creation_date})")
             return name
     return None
+
+def get_quarantined_stacks(region, prefix=None):
+    quarantined_stacks = []
+    cfn_client = boto3.client("cloudformation", region_name=region)
+
+    for stack in cfn_client.describe_stacks().get("Stacks", []):
+        stack_name = stack.get("StackName")
+        if not stack_name:
+            continue
+        if prefix and not stack_name.startswith(prefix):
+            continue
+        if any(tag.get("Key") == QUARANTINE_TAG_KEY for tag in stack.get("Tags", [])):
+            quarantined_stacks.append(stack_name)
+    return quarantined_stacks
+
+def is_quarantined_stack(region, stack_name):
+    return stack_name in get_quarantined_stacks(region)
+
+def quarantine_stacks(region, stack_names):
+    for stack_name in stack_names:
+        add_tags_to_stack(region, stack_name, QUARANTINE_TAGS)
+
+def add_tags_to_stack(region, stack_name, tags):
+    cfn = boto3.resource("cloudformation", region_name=region)
+    stack = cfn.Stack(stack_name)
+
+    stack.update(
+        UsePreviousTemplate=True,
+        Parameters=get_unchanged_stack_parameters(stack),
+        Capabilities=stack.capabilities,
+        DisableRollback=True,
+        Tags=tags,
+    )
+
+def get_unchanged_stack_parameters(stack):
+    return [
+        {
+            'ParameterKey': current_parameter.get('ParameterKey'),
+            'UsePreviousValue': True,
+        } for current_parameter in stack.parameters
+    ]