Add bucket name in batcher event format

airbnb · Dec 13, 2017 · 0159d33 · 0159d33
1 parent 2b65df8
commit 0159d33
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 8 deletions.
diff --git a/lambda_functions/batcher/main.py b/lambda_functions/batcher/main.py
@@ -52,7 +52,15 @@ def sqs_entry(self) -> Dict[str, str]:
         return {
             'Id': str(self._id),
             'MessageBody': json.dumps({
-                'Records': [{'s3': {'object': {'key': key}}} for key in self._keys]
+                'Records': [
+                    {
+                        's3': {
+                            'bucket': {'name': os.environ['S3_BUCKET_NAME']},
+                            'object': {'key': key}
+                        }
+                    }
+                    for key in self._keys
+                ]
             })
         }
 

diff --git a/tests/lambda_functions/batcher/main_test.py b/tests/lambda_functions/batcher/main_test.py
@@ -10,6 +10,7 @@
 from tests import common
 
 
+# TODO: mock.patch
 class MainTest(unittest.TestCase):
     """Test the batcher enqueuing everything from S3 into SQS."""
 
@@ -65,7 +66,12 @@ def test_batcher_one_object(self):
                     'Id': '0',
                     'MessageBody': json.dumps({
                         'Records': [
-                            {'s3': {'object': {'key': 'test-key-1'}}}
+                            {
+                                's3': {
+                                    'bucket': {'name': 'test_s3_bucket'},
+                                    'object': {'key': 'test-key-1'}
+                                }
+                            }
                         ]
                     })
                 }
@@ -103,8 +109,18 @@ def test_batcher_one_full_batch(self):
                     'Id': '0',
                     'MessageBody': json.dumps({
                         'Records': [
-                            {'s3': {'object': {'key': 'test-key-1'}}},
-                            {'s3': {'object': {'key': 'test-key-2'}}}
+                            {
+                                's3': {
+                                    'bucket': {'name': 'test_s3_bucket'},
+                                    'object': {'key': 'test-key-1'}
+                                }
+                            },
+                            {
+                                's3': {
+                                    'bucket': {'name': 'test_s3_bucket'},
+                                    'object': {'key': 'test-key-2'}
+                                }
+                            }
                         ]
                     })
                 }
@@ -154,16 +170,31 @@ def mock_list(**kwargs):
                     'Id': '0',
                     'MessageBody': json.dumps({
                         'Records': [
-                            {'s3': {'object': {'key': 'test-key-1'}}},
-                            {'s3': {'object': {'key': 'test-key-2'}}}
+                            {
+                                's3': {
+                                    'bucket': {'name': 'test_s3_bucket'},
+                                    'object': {'key': 'test-key-1'}
+                                }
+                            },
+                            {
+                                's3': {
+                                    'bucket': {'name': 'test_s3_bucket'},
+                                    'object': {'key': 'test-key-2'}
+                                }
+                            }
                         ]
                     })
                 },
                 {
                     'Id': '1',
                     'MessageBody': json.dumps({
                         'Records': [
-                            {'s3': {'object': {'key': 'test-key-3'}}}
+                            {
+                                's3': {
+                                    'bucket': {'name': 'test_s3_bucket'},
+                                    'object': {'key': 'test-key-3'}
+                                }
+                            }
                         ]
                     })
                 }
@@ -219,7 +250,12 @@ def test_batcher_invoke_with_continuation(self):
                     'Id': '0',
                     'MessageBody': json.dumps({
                         'Records': [
-                            {'s3': {'object': {'key': 'test-continuation-token'}}}
+                            {
+                                's3': {
+                                    'bucket': {'name': 'test_s3_bucket'},
+                                    'object': {'key': 'test-continuation-token'}
+                                }
+                            }
                         ]
                     })
                 }