diff --git a/chalice/test.py b/chalice/test.py index 84e2814f6..f3ce3a70c 100644 --- a/chalice/test.py +++ b/chalice/test.py @@ -4,7 +4,7 @@ import contextlib from types import TracebackType -from typing import Optional, Type, Generator, Dict, Any, List # noqa +from typing import Optional, Type, Generator, Dict, Any, List, Tuple # noqa from chalice import Chalice # noqa from chalice.config import Config @@ -16,6 +16,10 @@ class FunctionNotFoundError(Exception): pass +class InvalidDynamoStreamViewType(Exception): + pass + + class Client(object): def __init__(self, app, stage_name='dev', project_dir='.'): # type: (Chalice, str, str) -> None @@ -323,6 +327,44 @@ def generate_kinesis_event(self, message_bodies, } for body in message_bodies] return {'Records': records} + def generate_dynamodb_event(self, images, view_type="NEW_AND_OLD_IMAGES"): + # type: (List[Tuple[Dict, Dict]], str) -> Dict[str, Any] + # KEYS_ONLY/NEW_IMAGE/OLD_IMAGE / NEW_AND_OLD_IMAGES + def construct_record(old_image, new_image): + record = { + "dynamodb": { + "ApproximateCreationDateTime": 1545084650.987, + "SequenceNumber": "12345", + "SizeBytes": 12345, + "StreamViewType": view_type, + }, + "awsRegion": "us-west-2", + "eventID": "da037887f71a88a1f6f4cfd149709d5a", + "eventName": "INSERT", + "eventSource": "aws:dynamodb", + "eventSourceARN": ( + "arn:aws:dynamodb:us-west-2:12345:table/MyTable/stream/" + "2015-05-11T21:21:33.291" + ) + } + + if view_type == "KEYS_ONLY": + record["dynamodb"]["Keys"] = list(new_image.keys()) + elif view_type == "NEW_IMAGE": + record["dynamodb"]["NewImage"] = new_image + elif view_type == "OLD_IMAGE": + record["dynamodb"]["OldImage"] = old_image + elif view_type == "NEW_AND_OLD_IMAGES": + record["dynamodb"]["NewImage"] = new_image + record["dynamodb"]["OldImage"] = old_image + else: + raise InvalidDynamoStreamViewType(view_type) + + return record + + records = [construct_record(old_image, new_image) for old_image, new_image in images] + return {'Records': records} + class TestLambdaClient(BaseClient): def __init__(self, app, config): diff --git a/docs/source/api.rst b/docs/source/api.rst index 2dc153468..32b1b505b 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -1800,6 +1800,10 @@ Testing Generates a Kinesis event. + .. method:: generate_dynamodb_event(images, view_type='NEW_AND_OLD_IMAGES') + + Generates a DynamoDB Stream event. + .. class:: HTTPResponse() diff --git a/tests/unit/test_test.py b/tests/unit/test_test.py index e59bf092f..c48100949 100644 --- a/tests/unit/test_test.py +++ b/tests/unit/test_test.py @@ -313,6 +313,105 @@ def foo(event): assert response.payload == [b'foo', b'bar', b'baz'] +def test_can_generate_dynamodb_event(): + app = Chalice('dynamodb') + + @app.on_dynamodb_record( + stream_arn=('arn:aws:dynamodb:us-west-2:12345:table/MyTable/stream/' + '2015-05-11T21:21:33.291') + ) + def foo(event): + return [(record.old_image, record.new_image, record.keys) for record in event] + + old_image = {'PK': {'S': 'foo'}, 'SK': {'S': 'bar'}} + new_image = {'PK': {'S': 'hello'}, 'SK': {'S': 'world'}} + + with Client(app) as client: + event = client.events.generate_dynamodb_event( + images=[(old_image, new_image)] + ) + response = client.lambda_.invoke('foo', event) + assert len(response.payload) == 1 + assert response.payload[0][0] == old_image + assert response.payload[0][1] == new_image + assert response.payload[0][2] == None + + +def test_can_generate_dynamodb_event_keys_only(): + app = Chalice('dynamodb') + + @app.on_dynamodb_record( + stream_arn=('arn:aws:dynamodb:us-west-2:12345:table/MyTable/stream/' + '2015-05-11T21:21:33.291') + ) + def foo(event): + return [(record.old_image, record.new_image, record.keys) for record in event] + + old_image = {'PK': {'S': 'foo'}, 'SK': {'S': 'bar'}} + new_image = {'PK': {'S': 'hello'}, 'SK': {'S': 'world'}} + + with Client(app) as client: + event = client.events.generate_dynamodb_event( + images=[(old_image, new_image)], + view_type="KEYS_ONLY" + ) + response = client.lambda_.invoke('foo', event) + assert len(response.payload) == 1 + assert response.payload[0][0] == None + assert response.payload[0][1] == None + assert response.payload[0][2] == list(new_image.keys()) + + +def test_can_generate_dynamodb_event_old_image(): + app = Chalice('dynamodb') + + @app.on_dynamodb_record( + stream_arn=('arn:aws:dynamodb:us-west-2:12345:table/MyTable/stream/' + '2015-05-11T21:21:33.291') + ) + def foo(event): + return [(record.old_image, record.new_image, record.keys) for record in event] + + old_image = {'PK': {'S': 'foo'}, 'SK': {'S': 'bar'}} + new_image = {'PK': {'S': 'hello'}, 'SK': {'S': 'world'}} + + with Client(app) as client: + event = client.events.generate_dynamodb_event( + images=[(old_image, new_image)], + view_type="OLD_IMAGE" + ) + response = client.lambda_.invoke('foo', event) + assert len(response.payload) == 1 + assert response.payload[0][0] == old_image + assert response.payload[0][1] == None + assert response.payload[0][2] == None + + +def test_can_generate_dynamodb_event_new_image(): + app = Chalice('dynamodb') + + @app.on_dynamodb_record( + stream_arn=('arn:aws:dynamodb:us-west-2:12345:table/MyTable/stream/' + '2015-05-11T21:21:33.291') + ) + def foo(event): + return [(record.old_image, record.new_image, record.keys) for record in event] + + old_image = {'PK': {'S': 'foo'}, 'SK': {'S': 'bar'}} + new_image = {'PK': {'S': 'hello'}, 'SK': {'S': 'world'}} + + with Client(app) as client: + event = client.events.generate_dynamodb_event( + images=[(old_image, new_image)], + view_type="NEW_IMAGE" + ) + response = client.lambda_.invoke('foo', event) + assert len(response.payload) == 1 + assert response.payload[0][0] == None + assert response.payload[0][1] == new_image + assert response.payload[0][2] == None + + def test_can_mix_pure_lambda_and_event_handlers(): app = Chalice('lambda-only')