refactor(test): Pre-create MSK/MQ in companion stack for testing (aws#3909)

Author: vicheey

integration/combination/test_function_with_mq.py

@@ -5,7 +5,7 @@
 
 from integration.config.service_names import MQ
 from integration.helpers.base_test import BaseTest, nonblocking
-from integration.helpers.resource import current_region_does_not_support, generate_suffix
+from integration.helpers.resource import current_region_does_not_support
 
 
 @skipIf(current_region_does_not_support([MQ]), "MQ is not supported in this testing region")
@@ -16,32 +16,23 @@ def companion_stack_outputs(self, get_companion_stack_outputs):
 
     @parameterized.expand(
         [
-            ("combination/function_with_mq", "MQBrokerName", "MQBrokerUserSecretName", "PreCreatedSubnetOne"),
-            (
-                "combination/function_with_mq_using_autogen_role",
-                "MQBrokerName2",
-                "MQBrokerUserSecretName2",
-                "PreCreatedSubnetTwo",
-            ),
+            ("combination/function_with_mq",),
+            ("combination/function_with_mq_using_autogen_role",),
         ]
     )
     @nonblocking
-    def test_function_with_mq(self, file_name, mq_broker, mq_secret, subnet_key):
+    def test_function_with_mq(self, file_name):
         companion_stack_outputs = self.companion_stack_outputs
-        parameters = self.get_parameters(companion_stack_outputs, subnet_key)
-        secret_name = mq_secret + "-" + generate_suffix()
-        parameters.append(self.generate_parameter(mq_secret, secret_name))
-        secret_name = mq_broker + "-" + generate_suffix()
-        parameters.append(self.generate_parameter(mq_broker, secret_name))
+        parameters = [
+            self.generate_parameter("PreCreatedMqBrokerArn", companion_stack_outputs["PreCreatedMqBrokerArn"]),
+            self.generate_parameter(
+                "PreCreatedMqBrokerSecretArn", companion_stack_outputs["PreCreatedMqBrokerSecretArn"]
+            ),
+        ]
 
         self.create_and_verify_stack(file_name, parameters)
 
-        mq_client = self.client_provider.mq_client
-        mq_broker_id = self.get_physical_id_by_type("AWS::AmazonMQ::Broker")
-        broker_summary = get_broker_summary(mq_broker_id, mq_client)
-
-        self.assertEqual(len(broker_summary), 1, "One MQ cluster should be present")
-        mq_broker_arn = broker_summary[0]["BrokerArn"]
+        mq_broker_arn = companion_stack_outputs["PreCreatedMqBrokerArn"]
 
         lambda_client = self.client_provider.lambda_client
         function_name = self.get_physical_id_by_type("AWS::Lambda::Function")
@@ -54,15 +45,3 @@ def test_function_with_mq(self, file_name, mq_broker, mq_secret, subnet_key):
 
         self.assertEqual(event_source_mapping_function_arn, lambda_function_arn)
         self.assertEqual(event_source_mapping_mq_broker_arn, mq_broker_arn)
-
-    def get_parameters(self, dictionary, subnet_key):
-        parameters = []
-        parameters.append(self.generate_parameter("PreCreatedVpc", dictionary["PreCreatedVpc"]))
-        parameters.append(self.generate_parameter(subnet_key, dictionary[subnet_key]))
-        parameters.append(self.generate_parameter("PreCreatedInternetGateway", dictionary["PreCreatedInternetGateway"]))
-        return parameters
-
-
-def get_broker_summary(mq_broker_id, mq_client):
-    broker_summaries = mq_client.list_brokers()["BrokerSummaries"]
-    return [broker_summary for broker_summary in broker_summaries if broker_summary["BrokerId"] == mq_broker_id]

integration/combination/test_function_with_msk.py

@@ -4,13 +4,9 @@
 
 from integration.config.service_names import MSK
 from integration.helpers.base_test import BaseTest, nonblocking
-from integration.helpers.resource import current_region_does_not_support, generate_suffix
+from integration.helpers.resource import current_region_does_not_support
 
 
-# Mark this test suite as nonblocking tests since MSK Cluster creation can take
-# up to 30 minutes according to https://docs.aws.amazon.com/msk/latest/developerguide/troubleshooting.html#troubleshooting-cluster-stuck
-# This would cause the test to fail due to MSK Cluster did not stablize.
-# We should investigate any other cause of failures.
 @skipIf(current_region_does_not_support([MSK]), "MSK is not supported in this testing region")
 @nonblocking
 class TestFunctionWithMsk(BaseTest):
@@ -19,38 +15,25 @@ def companion_stack_outputs(self, get_companion_stack_outputs):
         self.companion_stack_outputs = get_companion_stack_outputs
 
     def test_function_with_msk_trigger(self):
-        companion_stack_outputs = self.companion_stack_outputs
-        parameters = self.get_parameters(companion_stack_outputs)
-        cluster_name = "MskCluster-" + generate_suffix()
-        parameters.append(self.generate_parameter("MskClusterName", cluster_name))
+        parameters = self.get_parameters()
         self._common_validations_for_MSK("combination/function_with_msk", parameters)
 
     def test_function_with_msk_trigger_using_manage_policy(self):
-        companion_stack_outputs = self.companion_stack_outputs
-        parameters = self.get_parameters(companion_stack_outputs)
-        cluster_name = "MskCluster2-" + generate_suffix()
-        parameters.append(self.generate_parameter("MskClusterName2", cluster_name))
+        parameters = self.get_parameters()
         self._common_validations_for_MSK("combination/function_with_msk_using_managed_policy", parameters)
 
     def test_function_with_msk_trigger_and_s3_onfailure_events_destinations(self):
-        companion_stack_outputs = self.companion_stack_outputs
-        parameters = self.get_parameters(companion_stack_outputs)
-        cluster_name = "MskCluster3-" + generate_suffix()
-        parameters.append(self.generate_parameter("MskClusterName3", cluster_name))
+        parameters = self.get_parameters()
         self._common_validations_for_MSK(
             "combination/function_with_msk_trigger_and_s3_onfailure_events_destinations", parameters
         )
 
     def test_function_with_msk_trigger_and_premium_features(self):
-        companion_stack_outputs = self.companion_stack_outputs
-        parameters = self.get_parameters(companion_stack_outputs)
-        cluster_name = "MskCluster4-" + generate_suffix()
-        parameters.append(self.generate_parameter("MskClusterName4", cluster_name))
+        parameters = self.get_parameters()
         self._common_validations_for_MSK("combination/function_with_msk_trigger_and_premium_features", parameters)
         event_source_mapping_result = self._common_validations_for_MSK(
             "combination/function_with_msk_trigger_and_confluent_schema_registry", parameters
         )
-        # Verify error handling properties are correctly set
         self.assertTrue(event_source_mapping_result.get("BisectBatchOnFunctionError"))
         self.assertEqual(event_source_mapping_result.get("MaximumRecordAgeInSeconds"), 3600)
         self.assertEqual(event_source_mapping_result.get("MaximumRetryAttempts"), 3)
@@ -59,15 +42,8 @@ def test_function_with_msk_trigger_and_premium_features(self):
     def _common_validations_for_MSK(self, file_name, parameters):
         self.create_and_verify_stack(file_name, parameters)
 
-        kafka_client = self.client_provider.kafka_client
+        msk_cluster_arn = self.companion_stack_outputs["PreCreatedMskClusterArn"]
 
-        msk_cluster_id = self.get_physical_id_by_type("AWS::MSK::Cluster")
-        cluster_info_list = kafka_client.list_clusters()["ClusterInfoList"]
-        cluster_info = [x for x in cluster_info_list if x["ClusterArn"] == msk_cluster_id]
-
-        self.assertEqual(len(cluster_info), 1, "One MSK cluster should be present")
-
-        msk_cluster_arn = cluster_info[0]["ClusterArn"]
         lambda_client = self.client_provider.lambda_client
         function_name = self.get_physical_id_by_type("AWS::Lambda::Function")
         lambda_function_arn = lambda_client.get_function_configuration(FunctionName=function_name)["FunctionArn"]
@@ -82,8 +58,7 @@ def _common_validations_for_MSK(self, file_name, parameters):
         self.assertEqual(event_source_mapping_kafka_cluster_arn, msk_cluster_arn)
         return event_source_mapping_result
 
-    def get_parameters(self, dictionary):
-        parameters = []
-        parameters.append(self.generate_parameter("PreCreatedSubnetOne", dictionary["PreCreatedSubnetOne"]))
-        parameters.append(self.generate_parameter("PreCreatedSubnetTwo", dictionary["PreCreatedSubnetTwo"]))
-        return parameters
+    def get_parameters(self):
+        return [
+            self.generate_parameter("PreCreatedMskClusterArn", self.companion_stack_outputs["PreCreatedMskClusterArn"]),
+        ]

integration/conftest.py

@@ -1,11 +1,13 @@
 import logging
+import os
 from pathlib import Path
 
 import boto3
 import botocore
 import pytest
 from botocore.exceptions import ClientError
 
+from integration.config.service_names import MQ, MSK
 from integration.helpers.base_test import S3_BUCKET_PREFIX
 from integration.helpers.client_provider import ClientProvider
 from integration.helpers.deployer.exceptions.exceptions import S3DoesNotExistException, ThrottlingError
@@ -20,7 +22,7 @@
 
 LOG = logging.getLogger(__name__)
 
-COMPANION_STACK_NAME = "sam-integ-stack-companion"
+COMPANION_STACK_NAME = os.environ.get("COMPANION_STACK_NAME", "sam-integ-stack-companion")
 COMPANION_STACK_TEMPLATE = "companion-stack.yaml"
 SAR_APP_TEMPLATE = "example-sar-app.yaml"
 SAR_APP_NAME = "sam-integration-test-sar-app"
@@ -62,14 +64,29 @@ def clean_all_integ_buckets():
 
 @pytest.fixture()
 def setup_companion_stack_once(tmpdir_factory, get_prefix):
+    # When COMPANION_STACK_NAME is set via env var, the stack is pre-created by the test platform
+    if os.environ.get("COMPANION_STACK_NAME"):
+        return
     tests_integ_dir = Path(__file__).resolve().parents[1]
     template_folder = Path(tests_integ_dir, "integration", "setup")
     companion_stack_template_path = Path(template_folder, COMPANION_STACK_TEMPLATE)
     cfn_client = ClientProvider().cfn_client
     output_dir = tmpdir_factory.mktemp("data")
     stack_name = get_prefix + COMPANION_STACK_NAME
     companion_stack = Stack(stack_name, companion_stack_template_path, cfn_client, output_dir)
-    companion_stack.create_or_update(_stack_exists(stack_name))
+    parameters = _companion_stack_parameters()
+    companion_stack.create_or_update(_stack_exists(stack_name), parameters)
+
+
+def _companion_stack_parameters():
+    """Return companion stack parameters, disabling MSK/MQ in unsupported regions."""
+
+    params = []
+    if current_region_does_not_support([MSK]):
+        params.append({"ParameterKey": "CreateMskCluster", "ParameterValue": "false"})
+    if current_region_does_not_support([MQ]):
+        params.append({"ParameterKey": "CreateMqBroker", "ParameterValue": "false"})
+    return params
 
 
 @pytest.fixture()
@@ -155,6 +172,8 @@ def get_s3_uri(file_name, uri_type, bucket, region):
 
 @pytest.fixture()
 def delete_companion_stack_once(get_prefix):
+    if os.environ.get("COMPANION_STACK_NAME"):
+        return
     if not get_prefix:
         ClientProvider().cfn_client.delete_stack(StackName=COMPANION_STACK_NAME)
 
@@ -179,7 +198,11 @@ def get_stack_outputs(stack_description):
 
 @pytest.fixture()
 def get_companion_stack_outputs(get_prefix):
-    companion_stack_description = get_stack_description(get_prefix + COMPANION_STACK_NAME)
+    if os.environ.get("COMPANION_STACK_NAME"):
+        stack_name = COMPANION_STACK_NAME
+    else:
+        stack_name = get_prefix + COMPANION_STACK_NAME
+    companion_stack_description = get_stack_description(stack_name)
     return get_stack_outputs(companion_stack_description)
 
 

integration/helpers/base_test.py

@@ -93,7 +93,7 @@ def setUpClass(cls):
         cls.tests_integ_dir = Path(__file__).resolve().parents[1]
         cls.resources_dir = Path(cls.tests_integ_dir, "resources")
         cls.template_dir = Path(cls.resources_dir, "templates")
-        cls.output_dir = Path(cls.tests_integ_dir, "tmp" + "-" + generate_suffix())
+        cls.output_dir = Path("/tmp", "tmp-" + generate_suffix())
         cls.expected_dir = Path(cls.resources_dir, "expected")
         cls.code_dir = Path(cls.resources_dir, "code")
         cls.session = boto3.session.Session()

integration/helpers/resource.py

@@ -145,21 +145,26 @@ def _get_region():
     return region
 
 
+_TMP_CONFIG_DIR = Path("/tmp/integration/config")
+
+
 def read_test_config_file(filename):
     """Reads test config file and returns the contents"""
     tests_integ_dir = Path(__file__).resolve().parents[1]
     test_config_file_path = Path(tests_integ_dir, "config", filename)
+    tmp_path = Path(_TMP_CONFIG_DIR, filename)
+    if not test_config_file_path.is_file() and tmp_path.is_file():
+        test_config_file_path = tmp_path
     if not test_config_file_path.is_file():
         return {}
     test_config = load_yaml(str(test_config_file_path))
     return test_config
 
 
 def write_test_config_file_to_json(filename, input):
-    """Reads test config file and returns the contents"""
-    tests_integ_dir = Path(__file__).resolve().parents[1]
-    test_config_file_path = Path(tests_integ_dir, "config", filename)
-    with open(test_config_file_path, "w") as f:
+    """Writes test config file as JSON to /tmp for portability across environments."""
+    _TMP_CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    with open(Path(_TMP_CONFIG_DIR, filename), "w") as f:
         json.dump(input, f)
 
 

integration/helpers/stack.py

@@ -19,10 +19,10 @@ def __init__(self, stack_name, template_path, cfn_client, output_dir):
         self.stack_description = None
         self.stack_resources = None
 
-    def create_or_update(self, update):
+    def create_or_update(self, update, parameters=None):
         output_template_path = self._generate_output_file_path(self.template_path, self.output_dir)
         transform_template(self.template_path, output_template_path)
-        self._deploy_stack(output_template_path, update)
+        self._deploy_stack(output_template_path, update, parameters)
 
     def delete(self):
         self.cfn_client.delete_stack(StackName=self.stack_name)

integration/resources/expected/combination/function_with_mq.json

@@ -7,32 +7,8 @@
     "LogicalResourceId": "MyLambdaExecutionRole",
     "ResourceType": "AWS::IAM::Role"
   },
-  {
-    "LogicalResourceId": "PublicSubnetRouteTableAssociation",
-    "ResourceType": "AWS::EC2::SubnetRouteTableAssociation"
-  },
-  {
-    "LogicalResourceId": "MQSecurityGroup",
-    "ResourceType": "AWS::EC2::SecurityGroup"
-  },
-  {
-    "LogicalResourceId": "MyMqBroker",
-    "ResourceType": "AWS::AmazonMQ::Broker"
-  },
-  {
-    "LogicalResourceId": "RouteTable",
-    "ResourceType": "AWS::EC2::RouteTable"
-  },
-  {
-    "LogicalResourceId": "MQBrokerUserSecret",
-    "ResourceType": "AWS::SecretsManager::Secret"
-  },
   {
     "LogicalResourceId": "MyLambdaFunctionMyMqEvent",
     "ResourceType": "AWS::Lambda::EventSourceMapping"
-  },
-  {
-    "LogicalResourceId": "Route",
-    "ResourceType": "AWS::EC2::Route"
   }
 ]

integration/resources/expected/combination/function_with_mq_using_autogen_role.json

@@ -7,32 +7,8 @@
     "LogicalResourceId": "MyLambdaFunctionRole",
     "ResourceType": "AWS::IAM::Role"
   },
-  {
-    "LogicalResourceId": "PublicSubnetRouteTableAssociation",
-    "ResourceType": "AWS::EC2::SubnetRouteTableAssociation"
-  },
-  {
-    "LogicalResourceId": "MQSecurityGroup",
-    "ResourceType": "AWS::EC2::SecurityGroup"
-  },
-  {
-    "LogicalResourceId": "MyMqBroker",
-    "ResourceType": "AWS::AmazonMQ::Broker"
-  },
-  {
-    "LogicalResourceId": "RouteTable",
-    "ResourceType": "AWS::EC2::RouteTable"
-  },
-  {
-    "LogicalResourceId": "MQBrokerUserSecret",
-    "ResourceType": "AWS::SecretsManager::Secret"
-  },
   {
     "LogicalResourceId": "MyLambdaFunctionMyMqEvent",
     "ResourceType": "AWS::Lambda::EventSourceMapping"
-  },
-  {
-    "LogicalResourceId": "Route",
-    "ResourceType": "AWS::EC2::Route"
   }
 ]

integration/resources/expected/combination/function_with_msk.json

@@ -7,10 +7,6 @@
     "LogicalResourceId": "MyLambdaExecutionRole",
     "ResourceType": "AWS::IAM::Role"
   },
-  {
-    "LogicalResourceId": "MyMskCluster",
-    "ResourceType": "AWS::MSK::Cluster"
-  },
   {
     "LogicalResourceId": "MyMskStreamProcessorMyMskEvent",
     "ResourceType": "AWS::Lambda::EventSourceMapping"

integration/resources/expected/combination/function_with_msk_trigger_and_premium_features.json

@@ -7,10 +7,6 @@
     "LogicalResourceId": "MyLambdaExecutionRole",
     "ResourceType": "AWS::IAM::Role"
   },
-  {
-    "LogicalResourceId": "MyMskCluster",
-    "ResourceType": "AWS::MSK::Cluster"
-  },
   {
     "LogicalResourceId": "MyMskStreamProcessorMyMskEvent",
     "ResourceType": "AWS::Lambda::EventSourceMapping"