tobywf · Nov 25, 2019 · Nov 22, 2019 · Nov 22, 2019 · Nov 22, 2019 · Nov 23, 2019
diff --git a/src/cloudformation_cli_python_lib/scheduler.py b/src/cloudformation_cli_python_lib/scheduler.py
@@ -0,0 +1,58 @@
+import json
+import logging
+from datetime import datetime, timedelta
+from uuid import uuid4
+
+# boto3 doesn't have stub files
+from boto3 import Session  # type: ignore
+
+from botocore.exceptions import ClientError  # type: ignore
+
+from .utils import HandlerRequest, KitchenSinkEncoder
+
+LOG = logging.getLogger(__name__)
+
+
+class CloudWatchScheduler:
+    def __init__(self, boto3_session: Session):
+        self.client = boto3_session.client("events")
+
+    def reschedule_after_minutes(
+        self, function_arn: str, minutes_from_now: int, handler_request: HandlerRequest
+    ) -> None:
+        cron = self._min_to_cron(max(minutes_from_now, 1))
+        uuid = str(uuid4())
+        rule_name = f"reinvoke-handler-{uuid}"
+        target_id = f"reinvoke-target-{uuid}"
+        handler_request.requestContext["cloudWatchEventsRuleName"] = rule_name
+        handler_request.requestContext["cloudWatchEventsTargetId"] = target_id
+        json_request = json.dumps(handler_request.serialize(), cls=KitchenSinkEncoder)
+        LOG.info("Scheduling re-invoke at %s (%s)", cron, uuid)
+        self.client.put_rule(Name=rule_name, ScheduleExpression=cron, State="ENABLED")
+        self.client.put_targets(
+            Rule=rule_name,
+            Targets=[{"Id": target_id, "Arn": function_arn, "Input": json_request}],
+        )
+
+    def cleanup_cloudwatch_events(self, rule_name: str, target_id: str) -> None:
+        try:
+            if target_id and rule_name:
+                self.client.remove_targets(Rule=rule_name, Ids=[target_id])
+        except ClientError as e:
+            LOG.error(
+                "Error cleaning CloudWatchEvents Target (targetId=%s): %s", target_id, e
+            )
+        try:
+            if rule_name:
+                self.client.delete_rule(Name=rule_name, Force=True)
+        except ClientError as e:
+            LOG.error(
+                "Error cleaning CloudWatchEvents (ruleName=%s): %s", rule_name, str(e)
+            )
+
+    @staticmethod
+    def _min_to_cron(minutes: int) -> str:
+        schedule_time = datetime.now() + timedelta(minutes=minutes)
+        # add another minute, as per java implementation
+        schedule_time = schedule_time + timedelta(minutes=1)
+        return schedule_time.strftime("cron('%M %H %d %m ? %Y')")
diff --git a/src/cloudformation_cli_python_lib/utils.py b/src/cloudformation_cli_python_lib/utils.py
@@ -60,6 +60,13 @@ def deserialize(cls, json_data: MutableMapping[str, Any]) -> "RequestData":
                 setattr(req_data, key, Credentials(**creds))
         return req_data
 
+    def serialize(self) -> Mapping[str, Any]:
+        return {
+            key: value.__dict__.copy() if key.endswith("Credentials") else value
+            for key, value in self.__dict__.items()
+            if value is not None
+        }
+
 
 # pylint: disable=too-many-instance-attributes
 @dataclass
@@ -75,14 +82,21 @@ class HandlerRequest:
     requestData: RequestData
     stackId: str
     nextToken: Optional[str] = None
-    requestContext: Mapping[str, Any] = field(default_factory=dict)
+    requestContext: MutableMapping[str, Any] = field(default_factory=dict)
 
     @classmethod
     def deserialize(cls, json_data: MutableMapping[str, Any]) -> "HandlerRequest":
         event = HandlerRequest(**json_data)
         event.requestData = RequestData.deserialize(json_data.get("requestData", {}))
         return event
 
+    def serialize(self) -> Mapping[str, Any]:
+        return {
+            key: value.serialize() if key == "requestData" else value
+            for key, value in self.__dict__.items()
+            if value is not None
+        }
+
 
 @dataclass
 class UnmodelledRequest:

diff --git a/tests/lib/scheduler_test.py b/tests/lib/scheduler_test.py
@@ -0,0 +1,134 @@
+# pylint: disable=redefined-outer-name,protected-access
+from datetime import datetime
+from unittest.mock import Mock, patch
+from uuid import uuid4
+
+import pytest
+from cloudformation_cli_python_lib.scheduler import CloudWatchScheduler
+
+from botocore.exceptions import ClientError
+
+
+@pytest.fixture
+def mock_boto3_session():
+    return Mock("cloudformation_cli_python_lib.scheduler.Session", autospec=True)()
+
+
+@pytest.fixture
+def mock_handler_request():
+    mock_request = Mock(
+        "cloudformation_cli_python_lib.utils.HandlerRequest", autospec=True
+    )()
+    mock_request.requestContext = {}
+    mock_request.serialize.return_value = {}
+    return mock_request
+
+
+def test_instantiates_boto3_client(mock_boto3_session):
+    cw_scheduler = CloudWatchScheduler(boto3_session=mock_boto3_session)
+    mock_boto3_session.client.assert_called_once_with("events")
+    assert cw_scheduler.client == mock_boto3_session.client.return_value
+
+
+@patch(
+    "cloudformation_cli_python_lib.scheduler.CloudWatchScheduler._min_to_cron",
+    return_value="cron('30 16 21 11 ? 2019')",
+)
+def test_reschedule_after_minutes_zero(
+    mock_min_to_cron, mock_boto3_session, mock_handler_request
+):
+    cw_scheduler = CloudWatchScheduler(boto3_session=mock_boto3_session)
+
+    # if called with zero, should call cron with a 1
+    cw_scheduler.reschedule_after_minutes("arn:goes:here", 0, mock_handler_request)
+    mock_min_to_cron.assert_called_once_with(1)
+
+
+@patch(
+    "cloudformation_cli_python_lib.scheduler.CloudWatchScheduler._min_to_cron",
+    return_value="cron('30 16 21 11 ? 2019')",
+)
+def test_reschedule_after_minutes_not_zero(
+    mock_min_to_cron, mock_boto3_session, mock_handler_request
+):
+    cw_scheduler = CloudWatchScheduler(boto3_session=mock_boto3_session)
+    # if called with another number, should use that
+    cw_scheduler.reschedule_after_minutes("arn:goes:here", 2, mock_handler_request)
+    mock_min_to_cron.assert_called_once_with(2)
+
+
+@patch(
+    "cloudformation_cli_python_lib.scheduler.uuid4",
+    autospec=True,
+    return_value=str(uuid4()),
+)
+def test_reschedule_after_minutes_success(
+    mock_uuid4, mock_boto3_session, mock_handler_request
+):
+    cw_scheduler = CloudWatchScheduler(boto3_session=mock_boto3_session)
+    with patch(
+        "cloudformation_cli_python_lib.scheduler.CloudWatchScheduler._min_to_cron",
+        return_value="cron('30 16 21 11 ? 2019')",
+    ):
+        cw_scheduler.reschedule_after_minutes("arn:goes:here", 2, mock_handler_request)
+
+    # should have made appropriate calls to create the schedule
+    cw_scheduler.client.put_targets.assert_called_once_with(
+        Rule=f"reinvoke-handler-{mock_uuid4.return_value}",
+        Targets=[
+            {
+                "Id": f"reinvoke-target-{mock_uuid4.return_value}",
+                "Arn": "arn:goes:here",
+                "Input": "{}",
+            }
+        ],
+    )
+    cw_scheduler.client.put_rule.assert_called_once_with(
+        Name=f"reinvoke-handler-{mock_uuid4.return_value}",
+        ScheduleExpression="cron('30 16 21 11 ? 2019')",
+        State="ENABLED",
+    )
+
+
+@patch("cloudformation_cli_python_lib.scheduler.LOG", autospec=True)
+def test_cleanup_cloudwatch_events_empty(mock_logger, mock_boto3_session):
+    cw_scheduler = CloudWatchScheduler(boto3_session=mock_boto3_session)
+
+    # cleanup should silently pass if rule/target are empty
+    cw_scheduler.cleanup_cloudwatch_events("", "")
+    assert cw_scheduler.client.remove_targets.called is False
+    assert cw_scheduler.client.delete_rule.called is False
+    assert mock_logger.error.called is False
+
+
+@patch("cloudformation_cli_python_lib.scheduler.LOG", autospec=True)
+def test_cleanup_cloudwatch_events_success(mock_logger, mock_boto3_session):
+    cw_scheduler = CloudWatchScheduler(boto3_session=mock_boto3_session)
+
+    # when rule_name and target_id are provided we should call events client and not
+    # log errors if the deletion succeeds
+    cw_scheduler.cleanup_cloudwatch_events("rulename", "targetid")
+    assert mock_logger.error.called is False
+    cw_scheduler.client.remove_targets.assert_called_once()
+    cw_scheduler.client.delete_rule.assert_called_once()
+
+
+@patch("cloudformation_cli_python_lib.scheduler.LOG", autospec=True)
+def test_cleanup_cloudwatch_events_boto_error(mock_logger, mock_boto3_session):
+    cw_scheduler = CloudWatchScheduler(boto3_session=mock_boto3_session)
+
+    # cleanup should catch and log boto failures
+    err = ClientError(error_response={"Error": {"Code": "1"}}, operation_name="mock")
+    cw_scheduler.client.remove_targets.side_effect = err
+    cw_scheduler.client.delete_rule.side_effect = err
+    cw_scheduler.cleanup_cloudwatch_events("rulename", "targetid")
+    assert mock_logger.error.call_count == 2
+    cw_scheduler.client.remove_targets.assert_called_once()
+    cw_scheduler.client.delete_rule.assert_called_once()
+
+
+@patch("cloudformation_cli_python_lib.scheduler.datetime", autospec=True)
+def test__min_to_cron(mock_datetime):
+    mock_datetime.now.return_value = datetime.fromisoformat("2019-01-01 01:01:01")
+    cron = CloudWatchScheduler._min_to_cron(1)
+    assert cron == "cron('03 01 01 01 ? 2019')"
diff --git a/tests/lib/utils_test.py b/tests/lib/utils_test.py
@@ -1,7 +1,7 @@
 import json
 
 import pytest
-from cloudformation_cli_python_lib.utils import KitchenSinkEncoder
+from cloudformation_cli_python_lib.utils import HandlerRequest, KitchenSinkEncoder
 
 import hypothesis.strategies as s
 from hypothesis import given
@@ -42,3 +42,59 @@ class Unserializable:
 
     with pytest.raises(TypeError):
         json.dumps(Unserializable(), cls=KitchenSinkEncoder)
+
+
+def test_handler_request_serde_roundtrip():
+    payload = {
+        "awsAccountId": "123456789012",
+        "bearerToken": "123456",
+        "region": "us-east-1",
+        "action": "CREATE",
+        "responseEndpoint": "https://cloudformation.us-west-2.amazonaws.com",
+        "resourceType": "AWS::Test::TestModel",
+        "resourceTypeVersion": "1.0",
+        "nextToken": None,
+        "requestContext": {
+            "invocation": 2,
+            "callbackContext": {"contextPropertyA": "Value"},
+            "cloudWatchEventsRuleName": "reinvoke-handler-4754ac8a-623b-45fe-84bc-f5394"
+            "118a8be",
+            "cloudWatchEventsTargetId": "reinvoke-target-4754ac8a-623b-45fe-84bc-f53941"
+            "18a8be",
+        },
+        "requestData": {
+            "callerCredentials": None,
+            "platformCredentials": {
+                "accessKeyId": "32IEHAHFIAG538KYASAI",
+                "secretAccessKey": "0O2hop/5vllVHjbA8u52hK8rLcroZpnL5NPGOi66",
+                "sessionToken": "gqe6eIsFPHOlfhc3RKl5s5Y6Dy9PYvN1CEYsswz5TQUsE8WfHD6LPK"
+                "549euXm4Vn4INBY9nMJ1cJe2mxTYFdhWHSnkOQv2SHemal",
+            },
+            "providerCredentials": {
+                "accessKeyId": "HDI0745692Y45IUTYR78",
+                "secretAccessKey": "4976TUYVI234/5GW87ERYG823RF87GY9EIUH452I3",
+                "sessionToken": "842HYOFIQAEUDF78R8T7IU43HSADYGIFHBJSDHFA87SDF9PYvN1CEY"
+                "ASDUYFT5TQ97YASIHUDFAIUEYRISDKJHFAYSUDTFSDFADS",
+            },
+            "providerLogGroupName": "providerLoggingGroupName",
+            "logicalResourceId": "myBucket",
+            "resourceProperties": {},
+            "previousResourceProperties": None,
+            "systemTags": {"aws:cloudformation:stack-id": "SampleStack"},
+            "stackTags": {"tag1": "abc"},
+            "previousStackTags": {"tag1": "def"},
+        },
+        "stackId": "arn:aws:cloudformation:us-east-1:123456789012:stack/SampleStack/e72"
+        "2ae60-fe62-11e8-9a0e-0ae8cc519968",
+    }
+    ser = HandlerRequest.deserialize(payload).serialize()
+    # remove None values from payload
+    expected = {
+        k: {k: v for k, v in payload["requestData"].items() if v is not None}
+        if k == "requestData"
+        else v
+        for k, v in payload.items()
+        if v is not None
+    }
+
+    assert ser == expected