[integ-tests framework] Honor --ami-owner option when retrieving AMIs in integration tests

The --ami-owner option was ignored by `retrieve_latest_ami`, so tests always searched the default per-OS owners
(amazon/self). This prevented running `retrieve_latest_ami` against ParallelCluster AMIs built in a separate account.

Author: hanwen-cluster

tests/integration-tests/conftest.py

@@ -903,6 +903,7 @@ def inject_additional_config_settings(  # noqa C901
                 ami_type="pcluster",
                 architecture=architecture,
                 additional_filters=_build_private_os_ami_filters(request),
+                request=request,
             ),
             ("Image", "CustomAmi"),
         )

tests/integration-tests/tests/common/utils.py

@@ -167,17 +167,24 @@ def retrieve_latest_ami(
                 and not request.config.getoption("pcluster_git_ref")
                 and not request.config.getoption("cookbook_git_ref")
                 and not request.config.getoption("node_git_ref")
+                and not request.config.getoption("ami_owner")
                 and not allow_private_ami
                 and os not in PRIVATE_OSES
             ):  # If none of Git refs is provided, the test is running against released version.
                 # Then retrieve public pcluster AMIs
                 additional_filters.append({"Name": "is-public", "Values": ["true"]})
         else:
             ami_name = _get_ami_for_os(ami_type, os, architecture).get("name")
+        if request and request.config.getoption("ami_owner") and ami_type == "pcluster":
+            owners = [
+                request.config.getoption("ami_owner")
+            ]  # ami owner should take effect only when getting pcluster AMI
+        else:
+            owners = _get_ami_for_os(ami_type, os, architecture).get("owners")
         describe_images_args = {
             "Filters": [{"Name": "name", "Values": [ami_name]}, {"Name": "architecture", "Values": [architecture]}]
             + additional_filters,
-            "Owners": _get_ami_for_os(ami_type, os, architecture).get("owners"),
+            "Owners": owners,
             "IncludeDeprecated": _get_ami_for_os(ami_type, os, architecture).get("includeDeprecated", False),
         }
         logging.info("Retrieving AMI with DescribeImages arguments: %s" % describe_images_args)

tests/integration-tests/tests/createami/test_createami.py

@@ -576,7 +576,7 @@ def test_build_image_custom_components(
     bucket.upload_file(str(test_datadir / custom_script_file), "scripts/custom_script.sh")
 
     # Get ParallelCluster AMI as base AMI
-    base_ami = retrieve_latest_ami(region, os, ami_type="pcluster", architecture=architecture)
+    base_ami = retrieve_latest_ami(region, os, ami_type="pcluster", architecture=architecture, request=request)
 
     image_id = generate_stack_name(
         "integ-tests-build-image-custom-components", request.config.getoption("stackname_suffix")

tests/integration-tests/tests/iam/test_iam_image.py

@@ -29,10 +29,13 @@ def test_iam_roles(
     pcluster_config_reader,
     images_factory,
     test_datadir,
+    request,
 ):
     instance_profile, lambda_cleanup_role = _create_image_roles(create_roles_stack)
 
-    image = _build_image(images_factory, instance_profile, lambda_cleanup_role, os, pcluster_config_reader, region)
+    image = _build_image(
+        images_factory, instance_profile, lambda_cleanup_role, os, pcluster_config_reader, region, request
+    )
 
     cfn_client = boto3.client("cloudformation", region_name=region)
     ec2_client = boto3.client("ec2", region_name=region)
@@ -53,11 +56,11 @@ def test_iam_roles(
     _wait_build_image_complete(image)
 
 
-def _build_image(images_factory, instance_profile, lambda_cleanup_role, os, pcluster_config_reader, region):
+def _build_image(images_factory, instance_profile, lambda_cleanup_role, os, pcluster_config_reader, region, request):
     # Generate image ID
     image_id = generate_stack_name("integ-tests-build-image", "")
     # Get base AMI
-    base_ami = retrieve_latest_ami(region, os, ami_type="pcluster", architecture="x86_64")
+    base_ami = retrieve_latest_ami(region, os, ami_type="pcluster", architecture="x86_64", request=request)
     image_config = pcluster_config_reader(
         config_file="image.config.yaml",
         parent_image=base_ami,

tests/integration-tests/tests/schedulers/conftest.py

@@ -191,7 +191,9 @@ def slurm_dbd(request, database, region, os, vpc_stack_for_database, munge_key):
                 {
                     "ParameterKey": "AmiId",
                     "ParameterValue": (
-                        custom_ami if custom_ami else retrieve_latest_ami(region, os, ami_type="pcluster")
+                        custom_ami
+                        if custom_ami
+                        else retrieve_latest_ami(region, os, ami_type="pcluster", request=request)
                     ),
                 },
                 {"ParameterKey": "DBMSClientSG", "ParameterValue": database.cfn_outputs["DatabaseClientSecurityGroup"]},