tobywf · Jan 27, 2020 · Dec 5, 2019 · Jan 16, 2020
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -72,6 +72,7 @@ repos:
       --cov-report=html
       --cov="rpdk.python"
       --cov="cloudformation_cli_python_lib"
+      --durations=5
       "tests/"
     language: system
     # ignore all files, run on hard-coded modules instead

diff --git a/setup.cfg b/setup.cfg
@@ -17,7 +17,7 @@ max-complexity = 10
 max-line-length = 88
 select = C,E,F,W,B,B950
 # C812, C815, W503 clash with black, F723 false positive
-ignore = E501,C812,C815,W503,F723
+ignore = E501,C812,C815,C816,W503,F723
 
 [isort]
 line_length = 88
@@ -28,6 +28,7 @@ skip = env
 include_trailing_comma = true
 combine_as_imports = True
 force_grid_wrap = 0
+known_standard_library = dataclasses
 known_first_party = rpdk
 known_third_party = boto3,jinja2,cloudformation_cli_python_lib,pytest
 

diff --git a/src/cloudformation_cli_python_lib/callback.py b/src/cloudformation_cli_python_lib/callback.py
@@ -3,17 +3,15 @@
 from typing import Optional
 from uuid import uuid4
 
-# boto3 doesn't have stub files
-from boto3 import Session  # type: ignore
-
+from .boto3_proxy import SessionProxy
 from .interface import BaseResourceModel, HandlerErrorCode, OperationStatus
 from .utils import KitchenSinkEncoder
 
 LOG = logging.getLogger(__name__)
 
 
 def report_progress(  # pylint: disable=too-many-arguments
-    session: Session,
+    session: SessionProxy,
     bearer_token: str,
     error_code: Optional[HandlerErrorCode],
     operation_status: OperationStatus,

diff --git a/src/cloudformation_cli_python_lib/log_delivery.py b/src/cloudformation_cli_python_lib/log_delivery.py
@@ -1,31 +1,28 @@
 import logging
 import time
-from typing import Any, Mapping, Optional
+from typing import Any, Optional
 
-# boto3 doesn't have stub files
-import boto3  # type: ignore
+from .boto3_proxy import SessionProxy
+from .utils import HandlerRequest
 
 
 class ProviderFilter(logging.Filter):
-    PROVIDER = ""
+    def __init__(self, provider: str):
+        super().__init__()
+        self.provider = provider
 
     def filter(self, record: logging.LogRecord) -> bool:
-        return not record.name.startswith(self.PROVIDER)
+        return not record.name.startswith(self.provider)
 
 
 class ProviderLogHandler(logging.Handler):
     def __init__(
-        self,
-        group: str,
-        stream: str,
-        creds: Mapping[str, str],
-        *args: Any,
-        **kwargs: Any,
+        self, group: str, stream: str, session: SessionProxy, *args: Any, **kwargs: Any
     ):
         super(ProviderLogHandler, self).__init__(*args, **kwargs)
         self.group = group
         self.stream = stream.replace(":", "__")
-        self.client = boto3.client("logs", **creds)
+        self.client = session.client("logs")
         self.sequence_token = ""
 
     @classmethod
@@ -36,48 +33,27 @@ def _get_existing_logger(cls) -> Optional["ProviderLogHandler"]:
         return None
 
     @classmethod
-    def setup(cls, event_data: Mapping[str, Any]) -> None:
-        try:
-            log_creds = event_data["requestData"]["providerCredentials"]
-        except KeyError:
-            log_creds = {}
-        try:
-            log_group = event_data["requestData"]["providerLogGroupName"]
-        except KeyError:
-            log_group = ""
-        try:
-            stream_name = (
-                f'{event_data["stackId"]}/'
-                f'{event_data["requestData"]["logicalResourceId"]}'
-            )
-        except KeyError:
-            stream_name = f'{event_data["awsAccountId"]}-{event_data["region"]}'
+    def setup(
+        cls, request: HandlerRequest, provider_sess: Optional[SessionProxy]
+    ) -> None:
+        log_group = request.requestData.providerLogGroupName
+        if request.stackId and request.requestData.logicalResourceId:
+            stream_name = f"{request.stackId}/{request.requestData.logicalResourceId}"
+        else:
+            stream_name = f"{request.awsAccountId}-{request.region}"
 
         log_handler = cls._get_existing_logger()
-        if log_creds and log_group:
+        if provider_sess and log_group:
             if log_handler:
                 # This is a re-used lambda container, log handler is already setup, so
                 # we just refresh the client with new creds
-                log_handler.client = boto3.client(
-                    "logs",
-                    aws_access_key_id=log_creds["accessKeyId"],
-                    aws_secret_access_key=log_creds["secretAccessKey"],
-                    aws_session_token=log_creds["sessionToken"],
-                )
+                log_handler.client = provider_sess.client("logs")
                 return
             # filter provider messages from platform
-            ProviderFilter.PROVIDER = (
-                event_data["resourceType"].replace("::", "_").lower()
-            )
-            logging.getLogger().handlers[0].addFilter(ProviderFilter())
+            provider = request.resourceType.replace("::", "_").lower()
+            logging.getLogger().handlers[0].addFilter(ProviderFilter(provider))
             log_handler = cls(
-                group=log_group,
-                stream=stream_name,
-                creds={
-                    "aws_access_key_id": log_creds["accessKeyId"],
-                    "aws_secret_access_key": log_creds["secretAccessKey"],
-                    "aws_session_token": log_creds["sessionToken"],
-                },
+                group=log_group, stream=stream_name, session=provider_sess
             )
             # add log handler to root, so that provider gets plugin logs too
             logging.getLogger().addHandler(log_handler)

diff --git a/src/cloudformation_cli_python_lib/metrics.py b/src/cloudformation_cli_python_lib/metrics.py
@@ -1,12 +1,10 @@
 import datetime
 import logging
-from typing import Any, List, Mapping
-
-# boto3 doesn't have stub files
-from boto3.session import Session  # type: ignore
+from typing import Any, List, Mapping, Optional
 
 from botocore.exceptions import ClientError  # type: ignore
 
+from .boto3_proxy import SessionProxy
 from .interface import Action, MetricTypes, StandardUnit
 
 LOG = logging.getLogger(__name__)
@@ -19,11 +17,9 @@ def format_dimensions(dimensions: Mapping[str, str]) -> List[Mapping[str, str]]:
 
 
 class MetricPublisher:
-    def __init__(self, account_id: str, resource_type: str, session: Session) -> None:
-        suffix = resource_type.replace("::", "/")
-        self.namespace = f"{METRIC_NAMESPACE_ROOT}/{account_id}/{suffix}"
-        self.resource_type = resource_type
+    def __init__(self, session: SessionProxy, namespace: str) -> None:
         self.client = session.client("cloudwatch")
+        self.namespace = namespace
 
     def publish_metric(  # pylint: disable-msg=too-many-arguments
         self,
@@ -52,21 +48,29 @@ def publish_metric(  # pylint: disable-msg=too-many-arguments
 
 
 class MetricsPublisherProxy:
-    def __init__(self) -> None:
+    @staticmethod
+    def _make_namespace(account_id: str, resource_type: str) -> str:
+        suffix = resource_type.replace("::", "/")
+        return f"{METRIC_NAMESPACE_ROOT}/{account_id}/{suffix}"
+
+    def __init__(self, account_id: str, resource_type: str) -> None:
+        self.namespace = self._make_namespace(account_id, resource_type)
+        self.resource_type = resource_type
         self._publishers: List[MetricPublisher] = []
 
-    def add_metrics_publisher(self, publisher: MetricPublisher) -> None:
-        self._publishers.append(publisher)
+    def add_metrics_publisher(self, session: Optional[SessionProxy]) -> None:
+        if session:
+            self._publishers.append(MetricPublisher(session, self.namespace))
 
     def publish_exception_metric(
         self, timestamp: datetime.datetime, action: Action, error: Any
     ) -> None:
+        dimensions: Mapping[str, str] = {
+            "DimensionKeyActionType": action.name,
+            "DimensionKeyExceptionType": str(type(error)),
+            "DimensionKeyResourceType": self.resource_type,
+        }
         for publisher in self._publishers:
-            dimensions: Mapping[str, str] = {
-                "DimensionKeyActionType": action.name,
-                "DimensionKeyExceptionType": str(type(error)),
-                "DimensionKeyResourceType": publisher.resource_type,
-            }
             publisher.publish_metric(
                 metric_name=MetricTypes.HandlerException,
                 dimensions=dimensions,
@@ -78,11 +82,11 @@ def publish_exception_metric(
     def publish_invocation_metric(
         self, timestamp: datetime.datetime, action: Action
     ) -> None:
+        dimensions = {
+            "DimensionKeyActionType": action.name,
+            "DimensionKeyResourceType": self.resource_type,
+        }
         for publisher in self._publishers:
-            dimensions = {
-                "DimensionKeyActionType": action.name,
-                "DimensionKeyResourceType": publisher.resource_type,
-            }
             publisher.publish_metric(
                 metric_name=MetricTypes.HandlerInvocationCount,
                 dimensions=dimensions,
@@ -94,11 +98,11 @@ def publish_invocation_metric(
     def publish_duration_metric(
         self, timestamp: datetime.datetime, action: Action, milliseconds: float
     ) -> None:
+        dimensions = {
+            "DimensionKeyActionType": action.name,
+            "DimensionKeyResourceType": self.resource_type,
+        }
         for publisher in self._publishers:
-            dimensions = {
-                "DimensionKeyActionType": action.name,
-                "DimensionKeyResourceType": publisher.resource_type,
-            }
             publisher.publish_metric(
                 metric_name=MetricTypes.HandlerInvocationDuration,
                 dimensions=dimensions,
@@ -110,12 +114,12 @@ def publish_duration_metric(
     def publish_log_delivery_exception_metric(
         self, timestamp: datetime.datetime, error: Any
     ) -> None:
+        dimensions = {
+            "DimensionKeyActionType": "ProviderLogDelivery",
+            "DimensionKeyExceptionType": str(type(error)),
+            "DimensionKeyResourceType": self.resource_type,
+        }
         for publisher in self._publishers:
-            dimensions: Mapping[str, str] = {
-                "DimensionKeyActionType": "ProviderLogDelivery",
-                "DimensionKeyExceptionType": str(type(error)),
-                "DimensionKeyResourceType": publisher.resource_type,
-            }
             publisher.publish_metric(
                 metric_name=MetricTypes.HandlerException,
                 dimensions=dimensions,