Add new parameter --credential to integration tests runner

Integration tests can be now launched in new optin region, specifying the --credential parameter for each new region to test.
Credential parameter is a comma separated list in the format region,endpoint,ARN,externalId.

Signed-off-by: Luca Carrogu <carrogu@amazon.com>

tests/integration-tests/README.md

@@ -193,6 +193,27 @@ tests_outputs
             └── ...
 ```
 
+## Cross Account Integration Tests
+If you want to distribute integration tests across multiple accounts you can make use of the `--credential` flag. 
+This is useful to overcome restrictions related to account limits and be compliant with a multi-region, multi-account 
+setup.
+
+When the `--credential` flag is specified and STS assume-role call is made in order to fetch temporary credentials to 
+be used to run tests in a specific region. 
+
+The `--credential` flag is in the form `<region>,<endpoint_url>,<ARN>,<external_id>` and needs to be specified for each 
+region you want to use with an STS assumed role (that usually means for every region you want to have in a separate 
+account).
+
+ * `region` is the region you want to test with an assumed STS role (which is in the target account where you want to 
+ launch the integration tests)
+ * `endpoint_url` is the STS endpoint url of the main account to be called in order to assume the delegation role
+ * `ARN` is the ARN of the delegation role in the optin region account to be assumed by the main account
+ * `external_id` is the external ID of the delegation role  
+
+By default, the delegation role lifetime last for one hour. Mind that if you are planning to launch tests that last 
+more than one hour.
+
 ## Write Integration Tests
 
 All integration tests are defined in the `integration-tests/tests` directory.

tests/integration-tests/cfn_stacks_factory.py

@@ -14,7 +14,7 @@
 from botocore.exceptions import ClientError
 from retrying import retry
 
-from utils import retrieve_cfn_outputs
+from utils import retrieve_cfn_outputs, set_credentials, unset_credentials
 
 
 class CfnStack:
@@ -41,8 +41,9 @@ def cfn_outputs(self):
 class CfnStacksFactory:
     """Manage creation and deletion of CloudFormation stacks."""
 
-    def __init__(self):
+    def __init__(self, credentials):
         self.__created_stacks = {}
+        self.__credentials = credentials
 
     def create_stack(self, stack):
         """
@@ -52,23 +53,28 @@ def create_stack(self, stack):
         """
         name = stack.name
         region = stack.region
-        id = self.__get_stack_internal_id(name, region)
-        if id in self.__created_stacks:
-            raise ValueError("Stack {0} already exists in region {1}".format(name, region))
-
-        logging.info("Creating stack {0} in region {1}".format(name, region))
-        self.__created_stacks[id] = stack
         try:
-            cfn_client = boto3.client("cloudformation", region_name=region)
-            result = cfn_client.create_stack(StackName=name, TemplateBody=stack.template)
-            stack.cfn_stack_id = result["StackId"]
-            final_status = self.__wait_for_stack_creation(stack.cfn_stack_id, cfn_client)
-            self.__assert_stack_status(final_status, "CREATE_COMPLETE")
-        except Exception as e:
-            logging.error("Creation of stack {0} in region {1} failed with exception: {2}".format(name, region, e))
-            raise
+            set_credentials(region, self.__credentials)
+
+            id = self.__get_stack_internal_id(name, region)
+            if id in self.__created_stacks:
+                raise ValueError("Stack {0} already exists in region {1}".format(name, region))
+
+            logging.info("Creating stack {0} in region {1}".format(name, region))
+            self.__created_stacks[id] = stack
+            try:
+                cfn_client = boto3.client("cloudformation", region_name=region)
+                result = cfn_client.create_stack(StackName=name, TemplateBody=stack.template)
+                stack.cfn_stack_id = result["StackId"]
+                final_status = self.__wait_for_stack_creation(stack.cfn_stack_id, cfn_client)
+                self.__assert_stack_status(final_status, "CREATE_COMPLETE")
+            except Exception as e:
+                logging.error("Creation of stack {0} in region {1} failed with exception: {2}".format(name, region, e))
+                raise
 
-        logging.info("Stack {0} created successfully in region {1}".format(name, region))
+            logging.info("Stack {0} created successfully in region {1}".format(name, region))
+        finally:
+            unset_credentials()
 
     @retry(
         stop_max_attempt_number=10,
@@ -77,22 +83,29 @@ def create_stack(self, stack):
     )
     def delete_stack(self, name, region):
         """Destroy a created cfn stack."""
-        id = self.__get_stack_internal_id(name, region)
-        if id in self.__created_stacks:
-            logging.info("Destroying stack {0} in region {1}".format(name, region))
-            try:
-                stack = self.__created_stacks[id]
-                cfn_client = boto3.client("cloudformation", region_name=stack.region)
-                cfn_client.delete_stack(StackName=stack.name)
-                final_status = self.__wait_for_stack_deletion(stack.cfn_stack_id, cfn_client)
-                self.__assert_stack_status(final_status, "DELETE_COMPLETE")
-            except Exception as e:
-                logging.error("Deletion of stack {0} in region {1} failed with exception: {2}".format(name, region, e))
-                raise
-            del self.__created_stacks[id]
-            logging.info("Stack {0} deleted successfully in region {1}".format(name, region))
-        else:
-            logging.warning("Couldn't find stack with name {0} in region. Skipping deletion.".format(name, region))
+        try:
+            set_credentials(region, self.__credentials)
+
+            id = self.__get_stack_internal_id(name, region)
+            if id in self.__created_stacks:
+                logging.info("Destroying stack {0} in region {1}".format(name, region))
+                try:
+                    stack = self.__created_stacks[id]
+                    cfn_client = boto3.client("cloudformation", region_name=stack.region)
+                    cfn_client.delete_stack(StackName=stack.name)
+                    final_status = self.__wait_for_stack_deletion(stack.cfn_stack_id, cfn_client)
+                    self.__assert_stack_status(final_status, "DELETE_COMPLETE")
+                except Exception as e:
+                    logging.error(
+                        "Deletion of stack {0} in region {1} failed with exception: {2}".format(name, region, e)
+                    )
+                    raise
+                del self.__created_stacks[id]
+                logging.info("Stack {0} deleted successfully in region {1}".format(name, region))
+            else:
+                logging.warning("Couldn't find stack with name {0} in region. Skipping deletion.".format(name, region))
+        finally:
+            unset_credentials()
 
     def delete_all_stacks(self):
         """Destroy all created stacks."""

tests/integration-tests/conftest.py

@@ -35,13 +35,23 @@
     check_marker_skip_list,
 )
 from jinja2 import Environment, FileSystemLoader
-from utils import create_s3_bucket, delete_s3_bucket, random_alphanumeric, to_snake_case
+from utils import (
+    create_s3_bucket,
+    delete_s3_bucket,
+    random_alphanumeric,
+    set_credentials,
+    to_snake_case,
+    unset_credentials,
+)
 from vpc_builder import Gateways, SubnetConfig, VPCConfig, VPCTemplateBuilder
 
 
 def pytest_addoption(parser):
     """Register argparse-style options and ini-style config values, called once at the beginning of a test run."""
     parser.addoption("--regions", help="aws region where tests are executed", default=["us-east-1"], nargs="+")
+    parser.addoption(
+        "--credential", help="STS credential endpoint, in the format <region>,<endpoint>,<ARN>,<externalId>.", nargs="+"
+    )
     parser.addoption("--instances", help="aws instances under test", default=["c5.xlarge"], nargs="+")
     parser.addoption("--oss", help="OSs under test", default=["alinux"], nargs="+")
     parser.addoption("--schedulers", help="schedulers under test", default=["slurm"], nargs="+")
@@ -159,6 +169,7 @@ def _add_properties_to_report(item):
 
 
 @pytest.fixture(scope="class")
+@pytest.mark.usefixtures("setup_sts_credentials")
 def clusters_factory(request):
     """
     Define a fixture to manage the creation and destruction of clusters.
@@ -296,7 +307,7 @@ def _get_default_template_values(vpc_stacks, region, request):
 @pytest.fixture(scope="session")
 def cfn_stacks_factory(request):
     """Define a fixture to manage the creation and destruction of CloudFormation stacks."""
-    factory = CfnStacksFactory()
+    factory = CfnStacksFactory(request.config.getoption("credential"))
     yield factory
     if not request.config.getoption("no_delete"):
         factory.delete_all_stacks()
@@ -323,6 +334,14 @@ def cfn_stacks_factory(request):
 }
 
 
+@pytest.fixture(scope="class", autouse=True)
+def setup_sts_credentials(region, request):
+    """Setup environment for the integ tests"""
+    set_credentials(region, request.config.getoption("credential"))
+    yield
+    unset_credentials()
+
+
 @pytest.fixture(scope="session", autouse=True)
 def vpc_stacks(cfn_stacks_factory, request):
     """Create VPC used by integ tests in all configured regions."""

tests/integration-tests/test_runner.py

@@ -90,6 +90,13 @@ def _init_argparser():
     parser.add_argument(
         "-r", "--regions", help="AWS region where tests are executed.", default=TEST_DEFAULTS.get("regions"), nargs="+"
     )
+    parser.add_argument(
+        "--credential",
+        action="append",
+        help="STS credential endpoint, in the format <region>,<endpoint>,<ARN>,<externalId>. "
+        "Could be specified multiple times.",
+        required=False,
+    )
     parser.add_argument(
         "-i", "--instances", help="AWS instances under test.", default=TEST_DEFAULTS.get("instances"), nargs="+"
     )
@@ -238,6 +245,10 @@ def _get_pytest_args(args, regions, log_file, out_dir):
     pytest_args.extend(["--key-name", args.key_name])
     pytest_args.extend(["--key-path", args.key_path])
 
+    if args.credential:
+        pytest_args.append("--credential")
+        pytest_args.extend(args.credential)
+
     if args.retry_on_failures:
         # Rerun tests on failures for one more time after 60 seconds delay
         pytest_args.extend(["--reruns", "1", "--reruns-delay", "60"])

tests/integration-tests/utils.py

@@ -10,6 +10,7 @@
 # This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
 import logging
+import os
 import random
 import re
 import shlex
@@ -19,6 +20,8 @@
 import boto3
 from retrying import retry
 
+from assertpy import assert_that
+
 
 def retry_if_subprocess_error(exception):
     """Return True if we should retry (in this case when it's a CalledProcessError), False otherwise"""
@@ -116,3 +119,71 @@ def delete_s3_bucket(bucket_name, region):
         bucket.delete()
     except boto3.client("s3").exceptions.NoSuchBucket:
         pass
+
+
+def set_credentials(region, credential_arg):
+    """
+    Set credentials for boto3 clients and cli commands
+
+    :param region: region of the bucket
+    :param credential_arg: credential list
+    """
+    if credential_arg:
+        # credentials = dict { region1: (endpoint1, arn1, external_id1),
+        #                      region2: (endpoint2, arn2, external_id2),
+        #                      [...],
+        #                    }
+        credentials = {
+            region: (endpoint, arn, external_id)
+            for region, endpoint, arn, external_id in [
+                tuple(credential_tuple.strip().split(","))
+                for credential_tuple in credential_arg
+                if credential_tuple.strip()
+            ]
+        }
+
+        if region in credentials:
+            credential_endpoint, credential_arn, credential_external_id = credentials.get(region)
+            aws_credentials = _retrieve_sts_credential(
+                credential_endpoint, credential_arn, credential_external_id, region
+            )
+
+            # Set credential for all boto3 client
+            boto3.setup_default_session(
+                aws_access_key_id=aws_credentials["AccessKeyId"],
+                aws_secret_access_key=aws_credentials["SecretAccessKey"],
+                aws_session_token=aws_credentials["SessionToken"],
+            )
+
+            # Set credential for all cli command e.g. pcluster create
+            os.environ["AWS_ACCESS_KEY_ID"] = aws_credentials["AccessKeyId"]
+            os.environ["AWS_SECRET_ACCESS_KEY"] = aws_credentials["SecretAccessKey"]
+            os.environ["AWS_SESSION_TOKEN"] = aws_credentials["SessionToken"]
+
+
+def _retrieve_sts_credential(credential_endpoint, credential_arn, credential_external_id, region):
+    match = re.search(r"https://sts\.(.*?)\.", credential_endpoint)
+    endpoint_region = match.group(1)
+
+    assert_that(credential_endpoint and endpoint_region and credential_arn and credential_external_id).is_true()
+
+    sts = boto3.client("sts", region_name=endpoint_region, endpoint_url=credential_endpoint)
+    assumed_role_object = sts.assume_role(
+        RoleArn=credential_arn, ExternalId=credential_external_id, RoleSessionName=region + "_integration_tests_session"
+    )
+    aws_credentials = assumed_role_object["Credentials"]
+
+    return aws_credentials
+
+
+def unset_credentials():
+    """Unset credentials"""
+    # Unset credential for all boto3 client
+    boto3.setup_default_session()
+    # Unset credential for cli command e.g. pcluster create
+    if "AWS_ACCESS_KEY_ID" in os.environ:
+        del os.environ["AWS_ACCESS_KEY_ID"]
+    if "AWS_SECRET_ACCESS_KEY" in os.environ:
+        del os.environ["AWS_SECRET_ACCESS_KEY"]
+    if "AWS_SESSION_TOKEN" in os.environ:
+        del os.environ["AWS_SESSION_TOKEN"]