return response for kinesis put_records

Rio Goodman · Rio Goodman · commit eb770142c31a · 2016-12-23T11:56:29.000-08:00
diff --git a/localstack/mock/infra.py b/localstack/mock/infra.py
@@ -9,6 +9,7 @@
 import requests
 import json
 import boto3
+import random
 import __init__
 from localstack.utils.aws import aws_stack
 from localstack.utils import common
@@ -35,6 +36,8 @@
 INSTALL_DIR_ES = '%s/elasticsearch' % INSTALL_DIR_INFRA
 TMP_ARCHIVE_ES = '/tmp/localstack.es.zip'
 
+ERROR_PROBABILITY = 0.05
+
 # set up logger
 LOGGER = logging.getLogger(__name__)
 
@@ -365,7 +368,12 @@ def update_apigateway(method, path, data, headers, response=None, return_forward
         return True
 
 
-def update_kinesis(method, path, data, headers, response=None, return_forward_info=False):
+def update_kinesis(method, path, data, headers, response=None, return_forward_info=False, return_errors=0):
+    # return_errors tells this method whether to include errors in its response
+    # 0: no errors
+    # 1: sometimes error response (governed by ERROR_PROBABILITY)
+    # 2: always error response
+    return_errors = 1
     if return_forward_info:
         return True
 
@@ -388,6 +396,7 @@ def update_kinesis(method, path, data, headers, response=None, return_forward_in
             records.append(record)
         stream_name = data['StreamName']
         lambda_api.process_kinesis_records(records, stream_name)
+    return put_records_response(records, return_errors)
 
 
 def update_dynamodb(method, path, data, headers, response=None, return_forward_info=False):
@@ -470,6 +479,49 @@ def dynamodb_extract_keys(item, table_name):
     return result
 
 
+# helper methods for response object
+def put_records_response(records, return_errors):
+    response = {
+        'FailedRecordCount': 0,
+        'Records': [
+        ]
+    }
+    for rec in records:
+        response = mock_record(response, return_errors)
+    return response
+
+
+def mock_record(response, return_errors):
+    if return_errors == 0:
+        record = good_record()
+    elif return_errors == 1:
+        record = unreliable_record()
+    else:
+        record = error_record()
+    if 'ErrorCode' in record:
+        response['FailedRecordCount'] += 1
+    response['Records'].append(record)
+    return response
+
+
+def good_record():
+    return {"SequenceNumber": 1, "ShardId": 1}
+
+
+def error_record():
+    return {
+        "ErrorCode": "ProvisionedThroughputExceededException",
+        "ErrorMessage": "Rate exceeded for shard shardId-1 in stream X under account 1."
+    }
+
+
+def unreliable_record():
+    if (random.random() < ERROR_PROBABILITY):
+        return bad_record()
+    else:
+        return good_record()
+
+
 if __name__ == '__main__':
     print('Starting local dev environment. CTRL-C to quit.')
     logging.basicConfig(level=logging.WARNING)
