Add unit test for network interface setup

Author: hgreebe

cli/src/pcluster/templates/queues_stack.py

@@ -13,6 +13,7 @@
     DEFAULT_EPHEMERAL_DIR,
     NODE_BOOTSTRAP_TIMEOUT,
     OS_MAPPING,
+    P6E_GB200,
     PCLUSTER_COMPUTE_RESOURCE_NAME_TAG,
     PCLUSTER_QUEUE_NAME_TAG,
     PCLUSTER_S3_ARTIFACTS_DICT,
@@ -366,9 +367,8 @@ def add_network_interfaces(
     compute_resource,
     queue_lt_security_groups,
 ):
-    """Generate launch template network interfaces list"""
-
-    is_gb200 = compute_resource.instance_types[0] == "p6e-gb200.36xlarge"
+    """Generate launch template network interfaces list."""
+    is_gb200 = compute_resource.instance_types[0].split(".")[0] == P6E_GB200
     interface = "efa" if compute_resource.efa and compute_resource.efa.enabled and not is_gb200 else None
 
     compute_lt_nw_interfaces = [

cli/tests/pcluster/templates/test_queues_stack.py

@@ -1,11 +1,13 @@
 import json
+from unittest.mock import MagicMock
 
 import pytest
 from assertpy import assert_that
 from freezegun import freeze_time
 
 from pcluster.schemas.cluster_schema import ClusterSchema
 from pcluster.templates.cdk_builder import CDKTemplateBuilder
+from pcluster.templates.queues_stack import add_network_interfaces
 from pcluster.utils import load_json_dict, load_yaml_dict
 from tests.pcluster.aws.dummy_aws_api import mock_aws_api
 from tests.pcluster.models.dummy_s3_bucket import dummy_cluster_bucket, mock_bucket_object_utils
@@ -152,6 +154,77 @@ def test_compute_nodes_dna_json(
     assert_that(json.loads(compute_node_extra_json)).is_equal_to(expected_compute_node_extra_json)
 
 
+class NetworkCard:
+    def __init__(self, index, max_interfaces=1):
+        self._index = index
+        self._max_interfaces = max_interfaces
+
+    def network_card_index(self):
+        return self._index
+
+    def maximum_network_interfaces(self):
+        return self._max_interfaces
+
+
+@pytest.mark.parametrize(
+    "efa_enabled, instance_type, network_cards_list, expected_interfaces",
+    [
+        (
+            True,
+            "p6e-gb200.36xlarge",
+            [NetworkCard(0), NetworkCard(1), NetworkCard(2, 2), NetworkCard(3), NetworkCard(4, 2)],
+            [
+                {"network_card_index": 0, "interface_type": None, "device_index": 0},
+                {"network_card_index": 1, "interface_type": "efa-only", "device_index": 0},
+                {"network_card_index": 2, "interface_type": "efa", "device_index": 1},
+                {"network_card_index": 3, "interface_type": "efa-only", "device_index": 0},
+                {"network_card_index": 4, "interface_type": "efa", "device_index": 1},
+            ],
+        ),
+        (
+            False,
+            "p6e-gb200.36xlarge",
+            [NetworkCard(0), NetworkCard(1), NetworkCard(2, 2), NetworkCard(3), NetworkCard(4, 2)],
+            [
+                {"network_card_index": 0, "interface_type": None, "device_index": 0},
+                {"network_card_index": 2, "interface_type": None, "device_index": 1},
+                {"network_card_index": 4, "interface_type": None, "device_index": 1},
+            ],
+        ),
+        (
+            True,
+            "c6in.32xlarge",
+            [NetworkCard(0), NetworkCard(1, 2), NetworkCard(2, 2)],
+            [
+                {"network_card_index": 0, "interface_type": "efa", "device_index": 0},
+                {"network_card_index": 1, "interface_type": "efa", "device_index": 1},
+                {"network_card_index": 2, "interface_type": "efa", "device_index": 1},
+            ],
+        ),
+    ],
+)
+def test_add_compute_resource_launch_template(
+    mocker, efa_enabled, instance_type, test_datadir, network_cards_list, expected_interfaces
+):
+    mock_compute_resource = MagicMock()
+    mock_compute_resource.name = "test-compute-resource"
+    mock_compute_resource.instance_types = [instance_type]
+    mock_compute_resource.efa.enabled = efa_enabled
+    mock_compute_resource.network_cards_list = network_cards_list
+
+    mock_queue = MagicMock()
+    mock_queue.name = "test-queue"
+
+    network_interfaces = add_network_interfaces(mock_queue, mock_compute_resource, ["sg-12345"])
+
+    assert len(network_interfaces) == len(expected_interfaces)
+
+    for actual, expected in zip(network_interfaces, expected_interfaces):
+        assert actual.network_card_index == expected["network_card_index"]
+        assert actual.interface_type == expected["interface_type"]
+        assert actual.device_index == expected["device_index"]
+
+
 def render_join(elem: dict):
     sep = str(elem[0])
     body = elem[1]