From 00bf81ceebfafc773c90062a4bdd9a57313eba1c Mon Sep 17 00:00:00 2001 From: Alexander Mancevice Date: Sat, 29 May 2021 11:05:44 -0400 Subject: [PATCH] Add support for Step Functions --- main.tf | 8 +++++++- src/events.py | 4 ++++ src/index.py | 11 +++++++++-- src/slack.py | 17 ++++++++--------- src/states.py | 26 ++++++++++++++++++++++++++ tests/index_test.py | 39 +++++++++++++++++++++++++++++++-------- tests/slack_test.py | 2 +- tests/states_test.py | 30 ++++++++++++++++++++++++++++++ 8 files changed, 116 insertions(+), 21 deletions(-) create mode 100644 src/states.py create mode 100644 tests/states_test.py diff --git a/main.tf b/main.tf index a69b594..f7c424b 100644 --- a/main.tf +++ b/main.tf @@ -123,7 +123,7 @@ resource "aws_cloudwatch_event_rule" "post" { description = local.events.post.rule_description event_pattern = jsonencode({ - detail-type = [{ prefix = "api/" }] + detail-type = ["post"] source = [local.events.source] }) } @@ -254,6 +254,12 @@ data "aws_iam_policy_document" "inline" { resources = ["*"] } + + statement { + sid = "SendTaskStatus" + actions = ["states:SendTask*"] + resources = ["*"] + } } resource "aws_iam_role" "role" { diff --git a/src/events.py b/src/events.py index bc0719a..9d59468 100644 --- a/src/events.py +++ b/src/events.py @@ -23,6 +23,10 @@ def detail(self): def detail_type(self): return self.event['detail-type'] + @property + def task_token(self): + return self.detail.get('task-token') + class HttpEvent(Event): @property diff --git a/src/index.py b/src/index.py index 73044ce..ebddd99 100644 --- a/src/index.py +++ b/src/index.py @@ -7,6 +7,7 @@ from logger import logger from secrets import export from slack import Slack +from states import States export(SecretId=os.getenv('SECRET_ID')) EVENTS_BUS_NAME = os.getenv('EVENTS_BUS_NAME') @@ -35,6 +36,7 @@ token=SLACK_TOKEN, verify=not SLACK_DISABLE_VERIFICATION, ) +states = States() @slack.route('GET /health') @@ -130,8 +132,13 @@ def post_slash_cmd(event): @logger.bind def post(event, context=None): event = EventBridgeEvent(event) - result = slack.post(event.detail_type, **event.detail) - events.publish(f'result/{ event.detail_type }', result) + result = slack.post(**event.detail) + if result['ok']: + events.publish('result', result) + if result['ok'] and event.task_token: + states.succeed(event.task_token, json.dumps(result)) + elif event.task_token: + states.fail(event.task_token, result['error'], json.dumps(result)) return result diff --git a/src/slack.py b/src/slack.py index a14c511..9d45e73 100644 --- a/src/slack.py +++ b/src/slack.py @@ -86,28 +86,27 @@ def install_url(self): }) return urlunsplit(url + [query, fragment]) - def post(self, path, body=None, headers=None): + def post(self, url, body=None, headers=None, **_): # Prepare request data = body.encode('utf-8') headers = {k.lower(): v for k, v in (headers or {}).items()} # Execute request - url = f'https://slack.com/{ path }' logger.info('POST %s %s', url, body) req = Request(url=url, data=data, headers=headers, method='POST') res = urlopen(req) # Parse response resdata = res.read().decode() - ok = False - if res.headers['content-type'].startswith('application/json'): - resdata = json.loads(resdata) - ok = resdata['ok'] + try: + resjson = json.loads(resdata) + except Exception: # pragma: no cover + resjson = {'ok': False} # Log response & return - log = logger.info if ok else logger.error - log('RESPONSE [%d] %s', res.status, json.dumps(resdata)) - return resdata + log = logger.info if resjson['ok'] else logger.error + log('RESPONSE [%d] %s', res.status, resdata) + return resjson def randstate(self): chars = string.ascii_letters + '1234567890' diff --git a/src/states.py b/src/states.py new file mode 100644 index 0000000..b4d83a9 --- /dev/null +++ b/src/states.py @@ -0,0 +1,26 @@ +import json + +import boto3 + +from logger import logger + + +class States: + def __init__(self, boto3_session=None): + self.boto3_session = boto3_session or boto3.Session() + self.client = self.boto3_session.client('stepfunctions') + + def fail(self, task_token, error, cause): + params = dict(taskToken=task_token, error=error, cause=cause) + logger.info('SEND TASK FAILURE %s', json.dumps(params)) + return self.client.send_task_failure(**params) + + def heartbeat(self, task_token): + params = dict(taskToken=task_token) + logger.info('SEND TASK HEARTBEAT %s', json.dumps(params)) + return self.client.send_task_heartbeat(**params) + + def succeed(self, task_token, output): + params = dict(taskToken=task_token, output=output) + logger.info('SEND TASK SUCCESS %s', json.dumps(params)) + return self.client.send_task_success(**params) diff --git a/tests/index_test.py b/tests/index_test.py index ad46a4c..30f2430 100644 --- a/tests/index_test.py +++ b/tests/index_test.py @@ -16,6 +16,7 @@ def setup(self): index.events.publish = mock.MagicMock() index.slack.install = mock.MagicMock() + index.slack.post = mock.MagicMock() index.slack.install.return_value = ( {'ok': True}, 'https://example.com/success', @@ -242,17 +243,39 @@ def test_500(self): assert index.proxy(event) == exp def test_post(self): + index.slack.post.return_value = {'ok': True} event = { - 'detail-type': 'api/chat.postMessage', + 'detail-type': 'post', 'detail': { + 'url': 'https://slack.com/api/chat.postMessage', 'body': json.dumps({'text': 'FIZZ'}), - 'headers': {'content-type': 'application/json'} + 'headers': {'content-type': 'application/json; charset=utf-8'}, + 'task-token': '', }, } - index.slack.post = mock.MagicMock() index.post(event) - index.slack.post.assert_called_once_with( - 'api/chat.postMessage', - body=json.dumps({'text': 'FIZZ'}), - headers={'content-type': 'application/json'} - ) + index.slack.post.assert_called_once_with(**{ + 'url': 'https://slack.com/api/chat.postMessage', + 'body': json.dumps({'text': 'FIZZ'}), + 'headers': {'content-type': 'application/json; charset=utf-8'}, + 'task-token': '', + }) + + def test_post_fail(self): + index.slack.post.return_value = {'ok': False, 'error': 'fizz'} + event = { + 'detail-type': 'post', + 'detail': { + 'url': 'https://slack.com/api/chat.postMessage', + 'body': json.dumps({'text': 'FIZZ'}), + 'headers': {'content-type': 'application/json; charset=utf-8'}, + 'task-token': '', + }, + } + index.post(event) + index.slack.post.assert_called_once_with(**{ + 'url': 'https://slack.com/api/chat.postMessage', + 'body': json.dumps({'text': 'FIZZ'}), + 'headers': {'content-type': 'application/json; charset=utf-8'}, + 'task-token': '', + }) diff --git a/tests/slack_test.py b/tests/slack_test.py index b291087..2386d3e 100644 --- a/tests/slack_test.py +++ b/tests/slack_test.py @@ -97,7 +97,7 @@ def test_install_url(self, state, oauth_install_uri, exp): def test_post(self): ret = self.subject.post( - 'api/chat.postMessage', + 'https://slack.com/api/chat.postMessage', json.dumps({'text': 'FIZZ'}), {'content-type': 'application/json; charset=utf-8'}, ) diff --git a/tests/states_test.py b/tests/states_test.py new file mode 100644 index 0000000..8435739 --- /dev/null +++ b/tests/states_test.py @@ -0,0 +1,30 @@ +from unittest.mock import MagicMock + +from src.states import States + + +class TestStates: + def setup(self): + self.boto3_session = MagicMock() + self.subject = States(boto3_session=self.boto3_session) + + def test_fail(self): + self.subject.fail('', 'error', '{}') + self.subject.client.send_task_failure.assert_called_once_with( + taskToken='', + error='error', + cause='{}', + ) + + def test_heartbeat(self): + self.subject.heartbeat('') + self.subject.client.send_task_heartbeat.assert_called_once_with( + taskToken='', + ) + + def test_succeed(self): + self.subject.succeed('', {'fizz': 'buzz'}) + self.subject.client.send_task_success.assert_called_once_with( + taskToken='', + output={'fizz': 'buzz'}, + )