[integ-test] fix build image test

Signed-off-by: Hanwen <[email protected]>
hanwen-pcluste · Oct 30, 2024 · 402bc21 · 402bc21
1 parent 3d2d9d1
commit 402bc21
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 2 deletions.
diff --git a/tests/integration-tests/tests/common/utils.py b/tests/integration-tests/tests/common/utils.py
@@ -89,11 +89,24 @@
     "rocky9": {"name": "rocky9-hvm-*-*", "owners": PCLUSTER_AMI_OWNERS},
 }
 
+FIRST_STAGE_AMI_OWNERS = ["self", "447714826191"]
+OS_TO_FIRST_STAGE_AMI_NAME_MAP = {
+    "alinux2": {"name": "first-stage-aws-parallelcluster-*-amzn2-*", "owners": FIRST_STAGE_AMI_OWNERS},
+    "alinux2023": {"name": "first-stage-aws-parallelcluster-*-amzn2023-*", "owners": FIRST_STAGE_AMI_OWNERS},
+    "ubuntu2004": {"name": "first-stage-aws-parallelcluster-*-ubuntu-2004-*", "owners": FIRST_STAGE_AMI_OWNERS},
+    "ubuntu2204": {"name": "first-stage-aws-parallelcluster-*-ubuntu-2204-*", "owners": FIRST_STAGE_AMI_OWNERS},
+    "rhel8": {"name": "first-stage-aws-parallelcluster-*-rhel8-*", "owners": FIRST_STAGE_AMI_OWNERS},
+    "rocky8": {"name": "first-stage-aws-parallelcluster-*-rocky8-*", "owners": FIRST_STAGE_AMI_OWNERS},
+    "rhel9": {"name": "first-stage-aws-parallelcluster-*-rhel9-*", "owners": FIRST_STAGE_AMI_OWNERS},
+    "rocky9": {"name": "first-stage-aws-parallelcluster-*-rocky9-*", "owners": FIRST_STAGE_AMI_OWNERS},
+}
+
 AMI_TYPE_DICT = {
     "official": OS_TO_OFFICIAL_AMI_NAME_OWNER_MAP,
     "remarkable": OS_TO_REMARKABLE_AMI_NAME_OWNER_MAP,
     "pcluster": OS_TO_PCLUSTER_AMI_NAME_OWNER_MAP,
     "kernel4": OS_TO_KERNEL4_AMI_NAME_OWNER_MAP,
+    "first_stage": OS_TO_FIRST_STAGE_AMI_NAME_MAP
 }
 
 

diff --git a/tests/integration-tests/tests/createami/test_createami.py b/tests/integration-tests/tests/createami/test_createami.py
@@ -124,10 +124,12 @@ def test_build_image(
 
     # Get base AMI
     # remarkable AMIs are not available for ARM and ubuntu2204 yet
-    if os not in ["ubuntu2204", "alinux2023"]:
+    if os in ["alinux2", "ubuntu2004"]:
         base_ami = retrieve_latest_ami(region, os, ami_type="remarkable", architecture=architecture)
+    elif "rhel" in os or "rocky" in os:
+        base_ami = retrieve_latest_ami(region, os, ami_type="first_stage", architecture=architecture)
     else:
-        base_ami = retrieve_latest_ami(region, os, architecture=architecture)
+        base_ami = retrieve_latest_ami(region, os, ami_type="official", architecture=architecture)
 
     image_config = pcluster_config_reader(
         config_file="image.config.yaml", parent_image=base_ami, instance_role=instance_role, bucket_name=bucket_name

diff --git a/tests/integration-tests/tests/createami/test_createami/test_build_image/image.config.yaml b/tests/integration-tests/tests/createami/test_createami/test_build_image/image.config.yaml
@@ -14,8 +14,10 @@ Build:
     Tags:
         - Key: dummyBuildTag
           Value: dummyBuildTag
+{% if os in ["alinux2", "alinux2023"] %}
     UpdateOsPackages:
         Enabled: true
+{% endif %}
 
 CustomS3Bucket: {{ bucket_name }}