feat: Allow Boto3 Session to be configured (#1686)

Co-authored-by: Jacob Fuss <[email protected]>

Author: jfuss

samtranslator/sdk/parameter.py

@@ -1,6 +1,8 @@
 import boto3
 import copy
 
+from samtranslator.translator.arn_generator import ArnGenerator, NoRegionFound
+
 
 class SamParameterValues(object):
     """
@@ -58,21 +60,20 @@ def add_default_parameter_values(self, sam_template):
             if param_name not in self.parameter_values and isinstance(value, dict) and "Default" in value:
                 self.parameter_values[param_name] = value["Default"]
 
-    def add_pseudo_parameter_values(self):
+    def add_pseudo_parameter_values(self, session=None):
         """
         Add pseudo parameter values
         :return: parameter values that have pseudo parameter in it
         """
+
+        if session is None:
+            session = boto3.session.Session()
+
+        if not session.region_name:
+            raise NoRegionFound("AWS Region cannot be found")
+
         if "AWS::Region" not in self.parameter_values:
-            self.parameter_values["AWS::Region"] = boto3.session.Session().region_name
+            self.parameter_values["AWS::Region"] = session.region_name
 
         if "AWS::Partition" not in self.parameter_values:
-            region = boto3.session.Session().region_name
-
-            # neither boto nor botocore has any way of returning the partition value yet
-            if region.startswith("cn-"):
-                self.parameter_values["AWS::Partition"] = "aws-cn"
-            elif region.startswith("us-gov-"):
-                self.parameter_values["AWS::Partition"] = "aws-us-gov"
-            else:
-                self.parameter_values["AWS::Partition"] = "aws"
+            self.parameter_values["AWS::Partition"] = ArnGenerator.get_partition_name(session.region_name)

samtranslator/translator/arn_generator.py

@@ -1,7 +1,13 @@
 import boto3
 
 
+class NoRegionFound(Exception):
+    pass
+
+
 class ArnGenerator(object):
+    class_boto_session = None
+
     @classmethod
     def generate_arn(cls, partition, service, resource, include_account_id=True):
         if not service or not resource:
@@ -43,7 +49,17 @@ def get_partition_name(cls, region=None):
         if region is None:
             # Use Boto3 to get the region where code is running. This uses Boto's regular region resolution
             # mechanism, starting from AWS_DEFAULT_REGION environment variable.
-            region = boto3.session.Session().region_name
+
+            if ArnGenerator.class_boto_session is None:
+                region = boto3.session.Session().region_name
+            else:
+                region = ArnGenerator.class_boto_session.region_name
+
+        # If region is still None, then we could not find the region. This will only happen
+        # in the local context. When this is deployed, we will be able to find the region like
+        # we did before.
+        if region is None:
+            raise NoRegionFound("AWS Region cannot be found")
 
         # setting default partition to aws, this will be overwritten by checking the region below
         partition = "aws"

samtranslator/translator/translator.py

@@ -25,12 +25,13 @@
 from samtranslator.plugins.policies.policy_templates_plugin import PolicyTemplatesForResourcePlugin
 from samtranslator.policy_template_processor.processor import PolicyTemplatesProcessor
 from samtranslator.sdk.parameter import SamParameterValues
+from samtranslator.translator.arn_generator import ArnGenerator
 
 
 class Translator:
     """Translates SAM templates into CloudFormation templates"""
 
-    def __init__(self, managed_policy_map, sam_parser, plugins=None):
+    def __init__(self, managed_policy_map, sam_parser, plugins=None, boto_session=None):
         """
         :param dict managed_policy_map: Map of managed policy names to the ARNs
         :param sam_parser: Instance of a SAM Parser
@@ -41,6 +42,9 @@ def __init__(self, managed_policy_map, sam_parser, plugins=None):
         self.plugins = plugins
         self.sam_parser = sam_parser
         self.feature_toggle = None
+        self.boto_session = boto_session
+
+        ArnGenerator.class_boto_session = self.boto_session
 
     def _get_function_names(self, resource_dict, intrinsics_resolver):
         """
@@ -92,7 +96,7 @@ def translate(self, sam_template, parameter_values, feature_toggle=None):
         self.redeploy_restapi_parameters = dict()
         sam_parameter_values = SamParameterValues(parameter_values)
         sam_parameter_values.add_default_parameter_values(sam_template)
-        sam_parameter_values.add_pseudo_parameter_values()
+        sam_parameter_values.add_pseudo_parameter_values(self.boto_session)
         parameter_values = sam_parameter_values.parameter_values
         # Create & Install plugins
         sam_plugins = prepare_plugins(self.plugins, parameter_values)

tests/sdk/test_parameter.py

@@ -1,9 +1,10 @@
 from parameterized import parameterized, param
 
-import pytest
 from unittest import TestCase
 from samtranslator.sdk.parameter import SamParameterValues
-from mock import patch
+from mock import patch, Mock
+
+from samtranslator.translator.arn_generator import NoRegionFound
 
 
 class TestSAMParameterValues(TestCase):
@@ -101,3 +102,10 @@ def test_add_pseudo_parameter_values_aws_partition_not_override(self):
         sam_parameter_values = SamParameterValues(parameter_values)
         sam_parameter_values.add_pseudo_parameter_values()
         self.assertEqual(expected, sam_parameter_values.parameter_values)
+
+    def test_add_pseudo_parameter_values_raises_NoRegionFound(self):
+        boto_session_mock = Mock()
+        boto_session_mock.region_name = None
+        sam_parameter_values = SamParameterValues({})
+        with self.assertRaises(NoRegionFound):
+            sam_parameter_values.add_pseudo_parameter_values(session=boto_session_mock)

tests/translator/test_arn_generator.py

@@ -0,0 +1,35 @@
+from unittest import TestCase
+from parameterized import parameterized
+from mock import Mock, patch
+
+from samtranslator.translator.arn_generator import ArnGenerator, NoRegionFound
+
+
+class TestArnGenerator(TestCase):
+    def setUp(self):
+        ArnGenerator.class_boto_session = None
+
+    @parameterized.expand(
+        [("us-east-1", "aws"), ("cn-east-1", "aws-cn"), ("us-gov-west-1", "aws-us-gov"), ("US-EAST-1", "aws")]
+    )
+    def test_get_partition_name(self, region, expected):
+        actual = ArnGenerator.get_partition_name(region)
+
+        self.assertEqual(actual, expected)
+
+    @patch("boto3.session.Session.region_name", None)
+    def test_get_partition_name_raise_NoRegionFound(self):
+        with self.assertRaises(NoRegionFound):
+            ArnGenerator.get_partition_name(None)
+
+    def test_get_partition_name_from_boto_session(self):
+        boto_session_mock = Mock()
+        boto_session_mock.region_name = "us-east-1"
+
+        ArnGenerator.class_boto_session = boto_session_mock
+
+        actual = ArnGenerator.get_partition_name()
+
+        self.assertEqual(actual, "aws")
+
+        ArnGenerator.class_boto_session = None