feat(batch): Support for Lambda context access in batch processing (#1609)

* Added types and parameter for lambda context, added unit tests

* Refactor parameter checking

* Added test for malformed context handling

Author: erikayao93

Diff for: packages/batch/src/BasePartialProcessor.ts

@@ -3,6 +3,7 @@
  */
 import {
   BaseRecord,
+  BatchProcessingOptions,
   EventSourceDataClassTypes,
   FailureResponse,
   ResultType,
@@ -16,6 +17,8 @@ abstract class BasePartialProcessor {
 
   public handler: CallableFunction;
 
+  public options?: BatchProcessingOptions;
+
   public records: BaseRecord[];
 
   public successMessages: EventSourceDataClassTypes[];
@@ -92,11 +95,16 @@ abstract class BasePartialProcessor {
    */
   public register(
     records: BaseRecord[],
-    handler: CallableFunction
+    handler: CallableFunction,
+    options?: BatchProcessingOptions
   ): BasePartialProcessor {
     this.records = records;
     this.handler = handler;
 
+    if (options) {
+      this.options = options;
+    }
+
     return this;
   }
 

Diff for: packages/batch/src/BatchProcessor.ts

@@ -19,7 +19,8 @@ class BatchProcessor extends BasePartialBatchProcessor {
   ): Promise<SuccessResponse | FailureResponse> {
     try {
       const data = this.toBatchType(record, this.eventType);
-      const result = await this.handler(data);
+
+      const result = await this.handler(data, this.options);
 
       return this.successHandler(record, result);
     } catch (e) {

Diff for: packages/batch/src/processPartialResponse.ts

@@ -1,6 +1,7 @@
 import {
   BasePartialBatchProcessor,
   BaseRecord,
+  BatchProcessingOptions,
   EventType,
   PartialItemFailureResponse,
 } from '.';
@@ -15,7 +16,8 @@ import {
 const processPartialResponse = async (
   event: { Records: BaseRecord[] },
   recordHandler: CallableFunction,
-  processor: BasePartialBatchProcessor
+  processor: BasePartialBatchProcessor,
+  options?: BatchProcessingOptions
 ): Promise<PartialItemFailureResponse> => {
   if (!event.Records) {
     const eventTypes: string = Object.values(EventType).toString();
@@ -26,9 +28,8 @@ const processPartialResponse = async (
     );
   }
 
-  const records = event['Records'];
+  processor.register(event.Records, recordHandler, options);
 
-  processor.register(records, recordHandler);
   await processor.process();
 
   return processor.response();

Diff for: packages/batch/src/types.ts

@@ -1,9 +1,17 @@
 /**
  * Types for batch processing utility
  */
-import { DynamoDBRecord, KinesisStreamRecord, SQSRecord } from 'aws-lambda';
+import {
+  Context,
+  DynamoDBRecord,
+  KinesisStreamRecord,
+  SQSRecord,
+} from 'aws-lambda';
+
+type BatchProcessingOptions = {
+  context: Context;
+};
 
-// types from base.py
 type EventSourceDataClassTypes =
   | SQSRecord
   | KinesisStreamRecord
@@ -21,6 +29,7 @@ type PartialItemFailures = { itemIdentifier: string };
 type PartialItemFailureResponse = { batchItemFailures: PartialItemFailures[] };
 
 export type {
+  BatchProcessingOptions,
   BaseRecord,
   EventSourceDataClassTypes,
   ResultType,

Diff for: packages/batch/tests/helpers/handlers.ts

@@ -1,4 +1,5 @@
 import { DynamoDBRecord, KinesisStreamRecord, SQSRecord } from 'aws-lambda';
+import { BatchProcessingOptions } from '../../src';
 
 const sqsRecordHandler = (record: SQSRecord): string => {
   const body = record.body;
@@ -58,11 +59,29 @@ const asyncDynamodbRecordHandler = async (
   return body;
 };
 
+const handlerWithContext = (
+  record: SQSRecord,
+  options: BatchProcessingOptions
+): string => {
+  const context = options.context;
+
+  try {
+    if (context.getRemainingTimeInMillis() == 0) {
+      throw Error('No time remaining.');
+    }
+  } catch (e) {
+    throw Error('Context possibly malformed. Displaying context:\n' + context);
+  }
+
+  return record.body;
+};
+
 export {
   sqsRecordHandler,
   asyncSqsRecordHandler,
   kinesisRecordHandler,
   asyncKinesisRecordHandler,
   dynamodbRecordHandler,
   asyncDynamodbRecordHandler,
+  handlerWithContext,
 };

Diff for: packages/batch/tests/unit/BatchProcessor.test.ts

@@ -4,7 +4,12 @@
  * @group unit/batch/class/batchprocessor
  */
 
-import { BatchProcessingError, BatchProcessor, EventType } from '../../src';
+import {
+  BatchProcessingError,
+  BatchProcessingOptions,
+  BatchProcessor,
+  EventType,
+} from '../../src';
 import {
   sqsRecordFactory,
   kinesisRecordFactory,
@@ -17,10 +22,14 @@ import {
   asyncKinesisRecordHandler,
   dynamodbRecordHandler,
   asyncDynamodbRecordHandler,
+  handlerWithContext,
 } from '../../tests/helpers/handlers';
+import { helloworldContext as dummyContext } from '../../../commons/src/samples/resources/contexts';
+import { Context } from 'aws-lambda';
 
 describe('Class: BatchProcessor', () => {
   const ENVIRONMENT_VARIABLES = process.env;
+  const options: BatchProcessingOptions = { context: dummyContext };
 
   beforeEach(() => {
     jest.clearAllMocks();
@@ -418,4 +427,58 @@ describe('Class: BatchProcessor', () => {
       );
     });
   });
+
+  describe('Batch processing with Lambda context', () => {
+    test('Batch processing when context is provided and handler accepts', async () => {
+      // Prepare
+      const firstRecord = sqsRecordFactory('success');
+      const secondRecord = sqsRecordFactory('success');
+      const records = [firstRecord, secondRecord];
+      const processor = new BatchProcessor(EventType.SQS);
+
+      // Act
+      processor.register(records, handlerWithContext, options);
+      const processedMessages = await processor.process();
+
+      // Assess
+      expect(processedMessages).toStrictEqual([
+        ['success', firstRecord.body, firstRecord],
+        ['success', secondRecord.body, secondRecord],
+      ]);
+    });
+
+    test('Batch processing when context is provided and handler does not accept', async () => {
+      // Prepare
+      const firstRecord = sqsRecordFactory('success');
+      const secondRecord = sqsRecordFactory('success');
+      const records = [firstRecord, secondRecord];
+      const processor = new BatchProcessor(EventType.SQS);
+
+      // Act
+      processor.register(records, sqsRecordHandler, options);
+      const processedMessages = await processor.process();
+
+      // Assess
+      expect(processedMessages).toStrictEqual([
+        ['success', firstRecord.body, firstRecord],
+        ['success', secondRecord.body, secondRecord],
+      ]);
+    });
+
+    test('Batch processing when malformed context is provided and handler attempts to use', async () => {
+      // Prepare
+      const firstRecord = sqsRecordFactory('success');
+      const secondRecord = sqsRecordFactory('success');
+      const records = [firstRecord, secondRecord];
+      const processor = new BatchProcessor(EventType.SQS);
+      const badContext = { foo: 'bar' };
+      const badOptions = { context: badContext as unknown as Context };
+
+      // Act
+      processor.register(records, handlerWithContext, badOptions);
+      await expect(processor.process()).rejects.toThrowError(
+        BatchProcessingError
+      );
+    });
+  });
 });

Diff for: packages/batch/tests/unit/processPartialResponse.test.ts

@@ -11,6 +11,7 @@ import {
   SQSEvent,
 } from 'aws-lambda';
 import {
+  BatchProcessingOptions,
   BatchProcessor,
   EventType,
   PartialItemFailureResponse,
@@ -24,6 +25,7 @@ import {
 import {
   asyncSqsRecordHandler,
   dynamodbRecordHandler,
+  handlerWithContext,
   kinesisRecordHandler,
   sqsRecordHandler,
 } from '../../tests/helpers/handlers';
@@ -33,6 +35,7 @@ import { Custom as dummyEvent } from '../../../commons/src/samples/resources/eve
 describe('Function: processPartialResponse()', () => {
   const ENVIRONMENT_VARIABLES = process.env;
   const context = dummyContext;
+  const options: BatchProcessingOptions = { context: dummyContext };
 
   beforeEach(() => {
     jest.clearAllMocks();
@@ -84,6 +87,27 @@ describe('Function: processPartialResponse()', () => {
       // Assess
       expect(ret).toStrictEqual({ batchItemFailures: [] });
     });
+
+    test('Process partial response function call with context provided', async () => {
+      // Prepare
+      const records = [
+        sqsRecordFactory('success'),
+        sqsRecordFactory('success'),
+      ];
+      const batch = { Records: records };
+      const processor = new BatchProcessor(EventType.SQS);
+
+      // Act
+      const ret = await processPartialResponse(
+        batch,
+        handlerWithContext,
+        processor,
+        options
+      );
+
+      // Assess
+      expect(ret).toStrictEqual({ batchItemFailures: [] });
+    });
   });
 
   describe('Process partial response function call through handler', () => {
@@ -188,5 +212,35 @@ describe('Function: processPartialResponse()', () => {
         )
       );
     });
+
+    test('Process partial response through handler with context provided', async () => {
+      // Prepare
+      const records = [
+        sqsRecordFactory('success'),
+        sqsRecordFactory('success'),
+      ];
+      const processor = new BatchProcessor(EventType.SQS);
+      const event: SQSEvent = { Records: records };
+
+      const handler = async (
+        event: SQSEvent,
+        _context: Context
+      ): Promise<PartialItemFailureResponse> => {
+        const options: BatchProcessingOptions = { context: _context };
+
+        return await processPartialResponse(
+          event,
+          handlerWithContext,
+          processor,
+          options
+        );
+      };
+
+      // Act
+      const result = await handler(event, context);
+
+      // Assess
+      expect(result).toStrictEqual({ batchItemFailures: [] });
+    });
   });
 });