tobywf · Nov 22, 2019 · Nov 20, 2019 · Nov 20, 2019 · Nov 20, 2019 · Nov 20, 2019
diff --git a/.gitignore b/.gitignore
@@ -125,4 +125,8 @@ dmypy.json
 
 .DS_Store
 .idea/
+.vscode/
 *.tar.gz
+
+# We want to allow the tests/lib folder
+!tests/lib
diff --git a/src/cloudformation_cli_python_lib/interface.py b/src/cloudformation_cli_python_lib/interface.py
@@ -24,6 +24,17 @@ class Action(str, _AutoName):
     LIST = auto()
 
 
+class StandardUnit(str, _AutoName):
+    Count = auto()
+    Milliseconds = auto()
+
+
+class MetricTypes(str, _AutoName):
+    HandlerException = auto()
+    HandlerInvocationCount = auto()
+    HandlerInvocationDuration = auto()
+
+
 class OperationStatus(str, _AutoName):
     PENDING = auto()
     IN_PROGRESS = auto()

diff --git a/src/cloudformation_cli_python_lib/metrics.py b/src/cloudformation_cli_python_lib/metrics.py
@@ -0,0 +1,96 @@
+import datetime
+import logging
+from typing import Any, List, Mapping
+
+# boto3 doesn't have stub files
+from boto3.session import Session  # type: ignore
+
+from botocore.exceptions import ClientError  # type: ignore
+
+from .interface import Action, MetricTypes, StandardUnit
+
+LOG = logging.getLogger(__name__)
+
+
+def format_dimensions(dimensions: Mapping[str, str]) -> List[Mapping[str, str]]:
+    return [{"Name": key, "Value": value} for key, value in dimensions.items()]
+
+
+class MetricPublisher:
+    def __init__(self, namespace: str, session: Session) -> None:
+        self.namespace = namespace
+        self.client = session.client("cloudwatch")
+
+    def _publish_metric(  # pylint: disable-msg=too-many-arguments
+        self,
+        metric_name: MetricTypes,
+        dimensions: Mapping[str, str],
+        unit: StandardUnit,
+        value: float,
+        timestamp: datetime.datetime,
+    ) -> None:
+        try:
+            self.client.put_metric_data(
+                Namespace=self.namespace,
+                MetricData=[
+                    {
+                        "MetricName": metric_name,
+                        "Dimensions": format_dimensions(dimensions),
+                        "Unit": unit,
+                        "Timestamp": str(timestamp),
+                        "Value": value,
+                    }
+                ],
+            )
+
+        except ClientError as e:
+            LOG.error("An error occurred while publishing metrics: %s", str(e))
+
+    def publish_exception_metric(
+        self, timestamp: datetime.datetime, action: Action, error: Any
+    ) -> None:
+        dimensions: Mapping[str, str] = {
+            "DimensionKeyActionType": action,
+            "DimensionKeyExceptionType": str(type(error)),
+            "DimensionKeyResourceType": self.namespace,
+        }
+
+        self._publish_metric(
+            metric_name=MetricTypes.HandlerException,
+            dimensions=dimensions,
+            unit=StandardUnit.Count,
+            value=1.0,
+            timestamp=timestamp,
+        )
+
+    def publish_invocation_metric(
+        self, timestamp: datetime.datetime, action: Action
+    ) -> None:
+        dimensions = {
+            "DimensionKeyActionType": action,
+            "DimensionKeyResourceType": self.namespace,
+        }
+
+        self._publish_metric(
+            metric_name=MetricTypes.HandlerInvocationCount,
+            dimensions=dimensions,
+            unit=StandardUnit.Count,
+            value=1.0,
+            timestamp=timestamp,
+        )
+
+    def publish_duration_metric(
+        self, timestamp: datetime.datetime, action: Action, milliseconds: float
+    ) -> None:
+        dimensions = {
+            "DimensionKeyActionType": action,
+            "DimensionKeyResourceType": self.namespace,
+        }
+
+        self._publish_metric(
+            metric_name=MetricTypes.HandlerInvocationDuration,
+            dimensions=dimensions,
+            unit=StandardUnit.Milliseconds,
+            value=milliseconds,
+            timestamp=timestamp,
+        )
diff --git a/tests/lib/metrics_test.py b/tests/lib/metrics_test.py
@@ -0,0 +1,145 @@
+from datetime import datetime
+from unittest.mock import MagicMock, call, patch
+
+import boto3
+from cloudformation_cli_python_lib.interface import MetricTypes, StandardUnit
+from cloudformation_cli_python_lib.metrics import MetricPublisher, format_dimensions
+
+from botocore.stub import Stubber
+
+
+class MockSession:
+    def __init__(self, client):
+        self._client = client
+
+    def client(self, _name):
+        return self._client
+
+
+class MockLog:
+    def __init__(self, handler):
+        self.debug = handler
+
+
+def test_format_dimensions():
+    dimensions = {"MyDimensionKey": "val_1", "MyDimensionKey2": "val_2"}
+    result = format_dimensions(dimensions)
+    assert [
+        {"Name": "MyDimensionKey", "Value": "val_1"},
+        {"Name": "MyDimensionKey2", "Value": "val_2"},
+    ] == result
+
+
+@patch("cloudformation_cli_python_lib.metrics.LOG", auto_spec=True)
+def test_put_metric_catches_error(mock_logger):
+    client = boto3.client("cloudwatch")
+    stubber = Stubber(client)
+
+    stubber.add_client_error("put_metric_data", "InternalServiceError")
+    stubber.activate()
+
+    publisher = MetricPublisher("fake-namespace", MockSession(client))
+    fake_datetime = datetime(2019, 1, 1)
+    publisher.publish_exception_metric(fake_datetime, "CREATE", "fake-error")
+    stubber.deactivate()
+    expected_calls = [
+        call.error(
+            "An error occurred while publishing metrics: %s",
+            "An error occurred (InternalServiceError) when calling the "
+            "PutMetricData operation: ",
+        )
+    ]
+    assert expected_calls == mock_logger.mock_calls
+
+
+def test_publish_exception_metric():
+    mock_client = patch("boto3.client")
+    mock_client.return_value = MagicMock()
+
+    fake_datetime = datetime(2019, 1, 1)
+    publisher = MetricPublisher("fake-namespace", mock_client.return_value)
+    publisher.publish_exception_metric(
+        fake_datetime, "fake-action", Exception("fake-error")
+    )
+
+    expected_calls = [
+        call.client("cloudwatch"),
+        call.client().put_metric_data(
+            Namespace="fake-namespace",
+            MetricData=[
+                {
+                    "MetricName": MetricTypes.HandlerException,
+                    "Dimensions": [
+                        {"Name": "DimensionKeyActionType", "Value": "fake-action"},
+                        {
+                            "Name": "DimensionKeyExceptionType",
+                            "Value": "<class 'Exception'>",
+                        },
+                        {"Name": "DimensionKeyResourceType", "Value": "fake-namespace"},
+                    ],
+                    "Unit": StandardUnit.Count,
+                    "Timestamp": str(fake_datetime),
+                    "Value": 1.0,
+                }
+            ],
+        ),
+    ]
+    assert expected_calls == mock_client.return_value.mock_calls
+
+
+def test_publish_invocation_metric():
+    mock_client = patch("boto3.client")
+    mock_client.return_value = MagicMock()
+
+    fake_datetime = datetime(2019, 1, 1)
+    publisher = MetricPublisher("fake-namespace", mock_client.return_value)
+    publisher.publish_invocation_metric(fake_datetime, "fake-action")
+
+    expected_calls = [
+        call.client("cloudwatch"),
+        call.client().put_metric_data(
+            Namespace="fake-namespace",
+            MetricData=[
+                {
+                    "MetricName": MetricTypes.HandlerInvocationCount,
+                    "Dimensions": [
+                        {"Name": "DimensionKeyActionType", "Value": "fake-action"},
+                        {"Name": "DimensionKeyResourceType", "Value": "fake-namespace"},
+                    ],
+                    "Unit": StandardUnit.Count,
+                    "Timestamp": str(fake_datetime),
+                    "Value": 1.0,
+                }
+            ],
+        ),
+    ]
+    assert expected_calls == mock_client.return_value.mock_calls
+
+
+def test_publish_duration_metric():
+    mock_client = patch("boto3.client")
+    mock_client.return_value = MagicMock()
+
+    fake_datetime = datetime(2019, 1, 1)
+    publisher = MetricPublisher("fake-namespace", mock_client.return_value)
+    publisher.publish_duration_metric(fake_datetime, "fake-action", 100)
+
+    expected_calls = [
+        call.client("cloudwatch"),
+        call.client().put_metric_data(
+            Namespace="fake-namespace",
+            MetricData=[
+                {
+                    "MetricName": MetricTypes.HandlerInvocationDuration,
+                    "Dimensions": [
+                        {"Name": "DimensionKeyActionType", "Value": "fake-action"},
+                        {"Name": "DimensionKeyResourceType", "Value": "fake-namespace"},
+                    ],
+                    "Unit": StandardUnit.Milliseconds,
+                    "Timestamp": str(fake_datetime),
+                    "Value": 100,
+                }
+            ],
+        ),
+    ]
+    assert expected_calls == mock_client.return_value.mock_calls