[Test] Consolidated and extended the unit test to verify the configuration of the shared storage security group.

In particular, the unit tests verifies that the security group has the expected rules in the following cases:

  1. Case 1: All nodes use the default (managed) security group
  2. Case 2: All nodes use the same custom security group
  3. Case 3: All nodes use different custom security groups

Author: gmarciani

cli/tests/pcluster/templates/test_cluster_stack.py

@@ -1376,84 +1376,114 @@ def test_resource_combination_name(
     assert_that(combination_name).is_equal_to(expected_combination_name)
 
 
-def test_storage_security_group_deduplication(mocker, test_datadir):
-    """
-    Test that storage security group rules are deduplicated when head, compute, and login nodes share the same SG.
-
-    When head node, compute nodes, and login nodes all use the same security group (sg-12345678),
-    only one set of ingress/egress rules should be created for that security group, not three separate sets.
-    """
-    mock_aws_api(mocker)
-    mock_bucket(mocker)
-    mock_bucket_object_utils(mocker)
-
-    input_yaml = load_yaml_dict(test_datadir / "config-shared-sg.yaml")
-    cluster_config = ClusterSchema(cluster_name="clustername").load(input_yaml)
-
-    generated_template, _ = CDKTemplateBuilder().build_cluster_template(
-        cluster_config=cluster_config, bucket=dummy_cluster_bucket(), stack_name="clustername"
-    )
-
-    # The EFS storage security group must have 2 ingress rules:
-    #   * allow traffic from cluster nodes (port 2049)
-    #   * allow traffic from storage nodes (all traffic)
-    efs_sg_ingress_rules = [
-        (name, resource)
-        for name, resource in generated_template["Resources"].items()
-        if resource["Type"] == "AWS::EC2::SecurityGroupIngress" and name.startswith("EFS") and "SecurityGroup" in name
-    ]
-    assert_that(len(efs_sg_ingress_rules)).is_equal_to(2)
-
-    # The FSx Lustre storage security group must have 3 rules:
-    #   * allow traffic from cluster nodes (port 2049)
-    #   * allow traffic from cluster nodes (ports 1018-1023)
-    #   * allow traffic from storage nodes (all traffic)
-    fsx_sg_ingress_rules = [
-        (name, resource)
-        for name, resource in generated_template["Resources"].items()
-        if resource["Type"] == "AWS::EC2::SecurityGroupIngress" and name.startswith("FSX") and "SecurityGroup" in name
-    ]
-    assert_that(len(fsx_sg_ingress_rules)).is_equal_to(3)
-
-    # Verify each storage type has the expected unique source security groups
-    for _storage_type, rules in [("EFS", efs_sg_ingress_rules), ("FSX", fsx_sg_ingress_rules)]:
-        source_sgs = {
-            str(rule["Properties"].get("SourceSecurityGroupId"))
-            for _, rule in rules
-            if rule["Properties"].get("SourceSecurityGroupId")
-        }
-        # Should have 2 unique source SGs (shared SG + storage SG)
-        assert_that(len(source_sgs)).is_equal_to(2)
-
-
-def test_storage_security_group_port_restrictions(mocker, test_datadir):
+@pytest.mark.parametrize(
+    "head_sg, q1_sg, q2_sg, login1_sg, login2_sg, expected_efs_ingress_rules, expected_fsx_ingress_rules",
+    [
+        # Case 1: All nodes use the default (managed) security group
+        # EFS: rule(2049) * (head, compute, login1, login2) + rule(all) * (storage) = 5 rules
+        # FSx: rule(988,1018-1023) * (head, compute, login1, login2) + rule(all) * (storage) = 9 rules
+        (None, None, None, None, None, 5, 9),
+        # Case 2: All nodes use the same custom security group (deduplication)
+        # EFS: rule(2049) * (customSG) + rule(all) * (storage) = 2 rules
+        # FSx: rule(988,1018-1023) * (customSG) + rule(all) * (storage) = 3 rules
+        (
+            "sg-1234567891234567a",
+            "sg-1234567891234567a",
+            "sg-1234567891234567a",
+            "sg-1234567891234567a",
+            "sg-1234567891234567a",
+            2,
+            3,
+        ),
+        # Case 3: All nodes use different custom security groups
+        # EFS: rule(2049) * (head_sg, q1_sg, q2_sg, login1_sg, login2_sg) + rule(all) * (storage) = 6 rules
+        # FSx: rule(988,1018-1023) * (head_sg, q1_sg, q2_sg, login1_sg, login2_sg) + rule(all) * (storage) = 11 rules
+        (
+            "sg-1234567891234567a",
+            "sg-1234567891234567b",
+            "sg-1234567891234567c",
+            "sg-1234567891234567d",
+            "sg-1234567891234567e",
+            6,
+            11,
+        ),
+    ],
+    ids=[
+        "all_default_sg",
+        "all_same_custom_sg",
+        "all_different_custom_sg",
+    ],
+)
+def test_storage_security_group_port_restrictions(
+    mocker,
+    test_datadir,
+    pcluster_config_reader,
+    head_sg,
+    q1_sg,
+    q2_sg,
+    login1_sg,
+    login2_sg,
+    expected_efs_ingress_rules,
+    expected_fsx_ingress_rules,
+):
     """
     Test that storage security group rules use restricted ports for head/compute/login nodes.
 
     Security group rules should follow these principles:
     1. Storage-to-Storage: Allow all traffic (protocol -1)
     2. Head/Compute/Login nodes to EFS: Allow only TCP port 2049
     3. Head/Compute/Login nodes to FSx Lustre: Allow only TCP ports 988 and 1018-1023
+    4. Rules are deduplicated when nodes share the same security group
     """
     mock_aws_api(mocker)
     mock_bucket(mocker)
     mock_bucket_object_utils(mocker)
 
-    input_yaml = load_yaml_dict(test_datadir / "config-shared-sg.yaml")
-    cluster_config = ClusterSchema(cluster_name="clustername").load(input_yaml)
+    rendered_config_path = pcluster_config_reader(
+        "config.yaml",
+        head_sg=head_sg,
+        q1_sg=q1_sg,
+        q2_sg=q2_sg,
+        login1_sg=login1_sg,
+        login2_sg=login2_sg,
+    )
+
+    rendered_config = load_yaml_dict(rendered_config_path)
+    cluster_config = ClusterSchema(cluster_name="clustername").load(rendered_config)
 
     generated_template, _ = CDKTemplateBuilder().build_cluster_template(
         cluster_config=cluster_config, bucket=dummy_cluster_bucket(), stack_name="clustername"
     )
 
+    # Build expected source security groups based on config
+    # Custom SG is used directly as string, managed SG is referenced via {"Ref": "SgName"}
+    expected_source = set()
+    expected_source.add(head_sg) if head_sg else expected_source.add(("Ref", "HeadNodeSecurityGroup"))
+    expected_source.add(q1_sg) if q1_sg else expected_source.add(("Ref", "ComputeSecurityGroup"))
+    expected_source.add(q2_sg) if q2_sg else expected_source.add(("Ref", "ComputeSecurityGroup"))
+    expected_source.add(login1_sg) if login1_sg else expected_source.add(("Ref", "login1LoginNodesSecurityGroup"))
+    expected_source.add(login2_sg) if login2_sg else expected_source.add(("Ref", "login2LoginNodesSecurityGroup"))
+
+    def _normalize_source_sg(source_sg):
+        """Convert source SG to a comparable format."""
+        if isinstance(source_sg, str):
+            return source_sg
+        elif isinstance(source_sg, dict) and "Ref" in source_sg:
+            return ("Ref", source_sg["Ref"])
+        return source_sg
+
     # Test EFS storage - should only allow port 2049
-    efs_ingress_rules = [
-        (name, resource)
-        for name, resource in generated_template["Resources"].items()
-        if resource["Type"] == "AWS::EC2::SecurityGroupIngress" and name.startswith("EFS") and "SecurityGroup" in name
-    ]
+    efs_ingress_rules = get_resources(
+        generated_template,
+        type="AWS::EC2::SecurityGroupIngress",
+        name_regex=r"^EFS.*SecurityGroup",
+    )
+
+    # Verify rule count (deduplication check)
+    assert_that(len(efs_ingress_rules)).is_equal_to(expected_efs_ingress_rules)
 
-    for name, rule in efs_ingress_rules:
+    efs_source_sgs = set()
+    for name, rule in efs_ingress_rules.items():
         props = rule["Properties"]
         if "Storage" in name:
             # Storage-to-Storage: all traffic allowed
@@ -1465,31 +1495,47 @@ def test_storage_security_group_port_restrictions(mocker, test_datadir):
             assert_that(props["IpProtocol"]).is_equal_to("tcp")
             assert_that(props["FromPort"]).is_equal_to(2049)
             assert_that(props["ToPort"]).is_equal_to(2049)
+            # Collect source security group
+            source_sg = _normalize_source_sg(props.get("SourceSecurityGroupId"))
+            efs_source_sgs.add(source_sg)
+
+    # Verify source SGs match expected
+    assert_that(efs_source_sgs).is_equal_to(expected_source)
 
     # Test FSx Lustre storage - should only allow ports 988 and 1018-1023
-    fsx_ingress_rules = [
-        (name, resource)
-        for name, resource in generated_template["Resources"].items()
-        if resource["Type"] == "AWS::EC2::SecurityGroupIngress" and name.startswith("FSX") and "SecurityGroup" in name
-    ]
+    fsx_ingress_rules = get_resources(
+        generated_template,
+        type="AWS::EC2::SecurityGroupIngress",
+        name_regex=r"^FSX.*SecurityGroup",
+    )
+
+    # Verify rule count (deduplication check)
+    assert_that(len(fsx_ingress_rules)).is_equal_to(expected_fsx_ingress_rules)
 
     # Collect non-storage rules to verify FSx ports
-    fsx_node_rules = [(name, rule) for name, rule in fsx_ingress_rules if "Storage" not in name]
-    fsx_storage_rules = [(name, rule) for name, rule in fsx_ingress_rules if "Storage" in name]
+    fsx_node_rules = {name: rule for name, rule in fsx_ingress_rules.items() if "Storage" not in name}
+    fsx_storage_rules = {name: rule for name, rule in fsx_ingress_rules.items() if "Storage" in name}
 
     # Verify Storage-to-Storage rule allows all traffic
-    for _name, rule in fsx_storage_rules:
+    for _name, rule in fsx_storage_rules.items():
         props = rule["Properties"]
         assert_that(props["IpProtocol"]).is_equal_to("-1")
         assert_that(props["FromPort"]).is_equal_to(0)
         assert_that(props["ToPort"]).is_equal_to(65535)
 
     # Verify Head/Compute/Login rules use TCP and FSx Lustre ports (988, 1018-1023)
     fsx_ports_found = set()
-    for _name, rule in fsx_node_rules:
+    fsx_source_sgs = set()
+    for _name, rule in fsx_node_rules.items():
         props = rule["Properties"]
         assert_that(props["IpProtocol"]).is_equal_to("tcp")
         fsx_ports_found.add((props["FromPort"], props["ToPort"]))
+        # Collect source security group
+        source_sg = _normalize_source_sg(props.get("SourceSecurityGroupId"))
+        fsx_source_sgs.add(source_sg)
 
     # Should have rules for port 988 and port range 1018-1023
     assert_that(fsx_ports_found).contains((988, 988), (1018, 1023))
+
+    # Verify source SGs match expected
+    assert_that(fsx_source_sgs).is_equal_to(expected_source)