Use role credentials when constructing clients for contract tests and use regional sts endpoint when getting account id (#548)

* Use role credentials when constructing clients for contract tests and use regional sts endpoint when getting account id

* Don't use default credentials chain for boto helper get account

* Use black to format files

Author: nina-ctrlv

src/rpdk/core/boto_helpers.py

@@ -46,11 +46,16 @@ def get_temporary_credentials(session, key_names=BOTO_CRED_KEYS, role_arn=None):
             response = sts_client.assume_role(
                 RoleArn=role_arn, RoleSessionName=session_name
             )
-        except ClientError as e:
+        except ClientError:
+            # pylint: disable=W1201
             LOG.debug(
-                "Getting session token resulted in unknown ClientError", exc_info=e
+                "Getting session token resulted in unknown ClientError. "
+                + "Could not assume specified role '%s'.",
+                role_arn,
+            )
+            raise DownstreamError() from Exception(
+                "Could not assume specified role '{}'".format(role_arn)
             )
-            raise DownstreamError("Could not assume specified role") from e
         temp = response["Credentials"]
         creds = (temp["AccessKeyId"], temp["SecretAccessKey"], temp["SessionToken"])
     else:
@@ -78,7 +83,14 @@ def get_service_endpoint(service, region):
     return "https://" + endpoint_data["hostname"]
 
 
-def get_account(session):
-    sts_client = session.client("sts")
+def get_account(session, temporary_credentials):
+    sts_client = session.client(
+        "sts",
+        endpoint_url=get_service_endpoint("sts", session.region_name),
+        region_name=session.region_name,
+        aws_access_key_id=temporary_credentials["accessKeyId"],
+        aws_secret_access_key=temporary_credentials["secretAccessKey"],
+        aws_session_token=temporary_credentials["sessionToken"],
+    )
     response = sts_client.get_caller_identity()
     return response.get("Account")

src/rpdk/core/contract/resource_client.py

@@ -132,7 +132,7 @@ def __init__(
             self._session, LOWER_CAMEL_CRED_KEYS, role_arn
         )
         self.region = region
-        self.account = get_account(self._session)
+        self.account = get_account(self._session, self._creds)
         self.partition = self._get_partition()
         self._function_name = function_name
         if endpoint.startswith("http://"):

src/rpdk/core/contract/suite/handler_commons.py

@@ -50,7 +50,8 @@ def test_create_failure_if_repeat_writeable_id(resource_client, current_resource
 @response_contains_resource_model_equal_current_model
 def test_read_success(resource_client, current_resource_model):
     primay_identifier_only_model = create_model_with_properties_in_path(
-        current_resource_model.copy(), resource_client.primary_identifier_paths,
+        current_resource_model.copy(),
+        resource_client.primary_identifier_paths,
     )
     _status, response, _error_code = resource_client.call_and_assert(
         Action.READ, OperationStatus.SUCCESS, primay_identifier_only_model
@@ -61,7 +62,8 @@ def test_read_success(resource_client, current_resource_model):
 @failed_event(error_code=HandlerErrorCode.NotFound)
 def test_read_failure_not_found(resource_client, current_resource_model):
     primay_identifier_only_model = create_model_with_properties_in_path(
-        current_resource_model, resource_client.primary_identifier_paths,
+        current_resource_model,
+        resource_client.primary_identifier_paths,
     )
     _status, _response, error_code = resource_client.call_and_assert(
         Action.READ, OperationStatus.FAILED, primay_identifier_only_model
@@ -124,7 +126,8 @@ def test_update_failure_not_found(resource_client, current_resource_model):
 
 def test_delete_success(resource_client, current_resource_model):
     primay_identifier_only_model = create_model_with_properties_in_path(
-        current_resource_model, resource_client.primary_identifier_paths,
+        current_resource_model,
+        resource_client.primary_identifier_paths,
     )
     _status, response, _error_code = resource_client.call_and_assert(
         Action.DELETE, OperationStatus.SUCCESS, primay_identifier_only_model
@@ -135,7 +138,8 @@ def test_delete_success(resource_client, current_resource_model):
 @failed_event(error_code=HandlerErrorCode.NotFound)
 def test_delete_failure_not_found(resource_client, current_resource_model):
     primay_identifier_only_model = create_model_with_properties_in_path(
-        current_resource_model, resource_client.primary_identifier_paths,
+        current_resource_model,
+        resource_client.primary_identifier_paths,
     )
     _status, _response, error_code = resource_client.call_and_assert(
         Action.DELETE, OperationStatus.FAILED, primay_identifier_only_model

src/rpdk/core/contract/suite/handler_delete.py

@@ -36,7 +36,8 @@ def deleted_resource(resource_client):
         model = response["resourceModel"]
         test_input_equals_output(resource_client, input_model, model)
         primay_identifier_only_model = create_model_with_properties_in_path(
-            model, resource_client.primary_identifier_paths,
+            model,
+            resource_client.primary_identifier_paths,
         )
         _status, response, _error = resource_client.call_and_assert(
             Action.DELETE, OperationStatus.SUCCESS, primay_identifier_only_model

src/rpdk/core/contract/suite/handler_update_invalid.py

@@ -30,7 +30,8 @@ def contract_update_create_only_property(resource_client):
             assert response["message"]
         finally:
             primay_identifier_only_model = create_model_with_properties_in_path(
-                created_model, resource_client.primary_identifier_paths,
+                created_model,
+                resource_client.primary_identifier_paths,
             )
             resource_client.call_and_assert(
                 Action.DELETE, OperationStatus.SUCCESS, primay_identifier_only_model

src/rpdk/core/jsonutils/pointer.py

@@ -99,21 +99,21 @@ def fragment_decode(pointer, prefix="#", output=tuple):
 
 def fragment_list(segments, prefix="properties", output=list):
     """Decode all segments of a JSON pointer from the URI fragment
-        identifier representation.
+    identifier representation.
 
-        >>> fragment_list(["properties"])
-        []
-        >>> fragment_list(["properties", "foo", "bar"])
-        ['foo', 'bar']
-        >>> fragment_list(["properties", "foo", "bar"], output=tuple)
-        ('foo', 'bar')
-        >>> fragment_list(["properties", "0", "%20", "~0"])
-        ['0', ' ', '~']
-        >>> fragment_list(["foo"])
-        Traceback (most recent call last):
-        ...
-        ValueError: Expected prefix 'properties', but was 'foo'
-        """
+    >>> fragment_list(["properties"])
+    []
+    >>> fragment_list(["properties", "foo", "bar"])
+    ['foo', 'bar']
+    >>> fragment_list(["properties", "foo", "bar"], output=tuple)
+    ('foo', 'bar')
+    >>> fragment_list(["properties", "0", "%20", "~0"])
+    ['0', ' ', '~']
+    >>> fragment_list(["foo"])
+    Traceback (most recent call last):
+    ...
+    ValueError: Expected prefix 'properties', but was 'foo'
+    """
     decoded = (part_decode(unquote(segment)) for segment in segments)
     actual = next(decoded)
     if prefix != actual:

src/rpdk/core/test.py

@@ -17,7 +17,7 @@
 
 from rpdk.core.jsonutils.pointer import fragment_decode
 
-from .boto_helpers import create_sdk_session
+from .boto_helpers import create_sdk_session, get_temporary_credentials
 from .contract.contract_plugin import ContractPlugin
 from .contract.interface import Action
 from .contract.resource_client import ResourceClient
@@ -57,9 +57,12 @@ def temporary_ini_file():
         yield str(path)
 
 
-def get_cloudformation_exports(region_name, endpoint_url):
+def get_cloudformation_exports(region_name, endpoint_url, role_arn):
     session = create_sdk_session(region_name)
-    cfn_client = session.client("cloudformation", endpoint_url=endpoint_url)
+    temp_credentials = get_temporary_credentials(session, role_arn=role_arn)
+    cfn_client = session.client(
+        "cloudformation", endpoint_url=endpoint_url, **temp_credentials
+    )
     paginator = cfn_client.get_paginator("list_exports")
     pages = paginator.paginate()
     exports = {}
@@ -68,12 +71,12 @@ def get_cloudformation_exports(region_name, endpoint_url):
     return exports
 
 
-def render_jinja(overrides_string, region_name, endpoint_url):
+def render_jinja(overrides_string, region_name, endpoint_url, role_arn):
     env = Environment(autoescape=True)
     parsed_content = env.parse(overrides_string)
     variables = meta.find_undeclared_variables(parsed_content)
     if variables:
-        exports = get_cloudformation_exports(region_name, endpoint_url)
+        exports = get_cloudformation_exports(region_name, endpoint_url, role_arn)
         invalid_exports = variables - exports.keys()
         if len(invalid_exports) > 0:
             invalid_exports_message = (
@@ -89,14 +92,14 @@ def render_jinja(overrides_string, region_name, endpoint_url):
     return to_return
 
 
-def get_overrides(root, region_name, endpoint_url):
+def get_overrides(root, region_name, endpoint_url, role_arn):
     if not root:
         return empty_override()
 
     path = root / "overrides.json"
     try:
         with path.open("r", encoding="utf-8") as f:
-            overrides_raw = render_jinja(f.read(), region_name, endpoint_url)
+            overrides_raw = render_jinja(f.read(), region_name, endpoint_url, role_arn)
     except FileNotFoundError:
         LOG.debug("Override file '%s' not found. No overrides will be applied", path)
         return empty_override()
@@ -123,7 +126,7 @@ def get_overrides(root, region_name, endpoint_url):
 
 
 # pylint: disable=R0914
-def get_inputs(root, region_name, endpoint_url, value):
+def get_inputs(root, region_name, endpoint_url, value, role_arn):
     inputs = {}
     if not root:
         return None
@@ -144,7 +147,9 @@ def get_inputs(root, region_name, endpoint_url, value):
 
                 file_path = path / file
                 with file_path.open("r", encoding="utf-8") as f:
-                    overrides_raw = render_jinja(f.read(), region_name, endpoint_url)
+                    overrides_raw = render_jinja(
+                        f.read(), region_name, endpoint_url, role_arn
+                    )
                 overrides = {}
                 for pointer, obj in overrides_raw.items():
                     overrides[pointer] = obj
@@ -175,13 +180,17 @@ def test(args):
     project.load()
 
     overrides = get_overrides(
-        project.root, args.region, args.cloudformation_endpoint_url
+        project.root, args.region, args.cloudformation_endpoint_url, args.role_arn
     )
 
     index = 1
     while True:
         inputs = get_inputs(
-            project.root, args.region, args.cloudformation_endpoint_url, index
+            project.root,
+            args.region,
+            args.cloudformation_endpoint_url,
+            index,
+            args.role_arn,
         )
         if not inputs:
             break

tests/contract/test_resource_client.py

@@ -80,7 +80,7 @@ def resource_client():
 
     mock_sesh.client.assert_called_once_with("lambda", endpoint_url=endpoint)
     mock_creds.assert_called_once_with(mock_sesh, LOWER_CAMEL_CRED_KEYS, None)
-    mock_account.assert_called_once_with(mock_sesh)
+    mock_account.assert_called_once_with(mock_sesh, {})
     assert client._creds == {}
     assert client._function_name == DEFAULT_FUNCTION
     assert client._schema == {}
@@ -121,7 +121,7 @@ def resource_client_inputs():
 
     mock_sesh.client.assert_called_once_with("lambda", endpoint_url=endpoint)
     mock_creds.assert_called_once_with(mock_sesh, LOWER_CAMEL_CRED_KEYS, None)
-    mock_account.assert_called_once_with(mock_sesh)
+    mock_account.assert_called_once_with(mock_sesh, {})
 
     assert client._creds == {}
     assert client._function_name == DEFAULT_FUNCTION
@@ -213,7 +213,9 @@ def test_init_sam_cli_client():
         "rpdk.core.contract.resource_client.create_sdk_session", autospec=True
     )
     patch_creds = patch(
-        "rpdk.core.contract.resource_client.get_temporary_credentials", autospec=True
+        "rpdk.core.contract.resource_client.get_temporary_credentials",
+        autospec=True,
+        return_value={},
     )
     patch_account = patch(
         "rpdk.core.contract.resource_client.get_account",
@@ -232,7 +234,7 @@ def test_init_sam_cli_client():
         "lambda", endpoint_url=DEFAULT_ENDPOINT, use_ssl=False, verify=False, config=ANY
     )
     mock_creds.assert_called_once_with(mock_sesh, LOWER_CAMEL_CRED_KEYS, None)
-    mock_account.assert_called_once_with(mock_sesh)
+    mock_account.assert_called_once_with(mock_sesh, {})
     assert client.account == ACCOUNT
 
 

tests/test_boto_helpers.py

@@ -200,10 +200,26 @@ def test_get_temporary_credentials_assume_role():
     assert tuple(creds.values()) == (access_key, secret_key, token)
 
 
-def test_get_account():
+def test_get_account_with_temporary_credentials():
     session = create_autospec(spec=Session, spec_set=True)
     client = session.client.return_value
-    get_account(session)
+    session.region_name = "us-east-1"
+    access_key = object()
+    secret_key = object()
+    token = object()
+    creds = {
+        "accessKeyId": access_key,
+        "secretAccessKey": secret_key,
+        "sessionToken": token,
+    }
+    get_account(session, creds)
 
-    session.client.assert_called_once_with("sts")
+    session.client.assert_called_once_with(
+        "sts",
+        aws_access_key_id=access_key,
+        aws_secret_access_key=secret_key,
+        aws_session_token=token,
+        endpoint_url="https://sts.us-east-1.amazonaws.com",
+        region_name="us-east-1",
+    )
     client.get_caller_identity.assert_called_once()