implement state parameter for state machine steps

janis-commerce · Jan 18, 2024 · 26ed536 · 26ed536
1 parent 7ec33bf
commit 26ed536
Show file tree

Hide file tree

Showing 7 changed files with 93 additions and 54 deletions.
diff --git a/lib/handler.js b/lib/handler.js
@@ -34,14 +34,15 @@ module.exports = class Handler {
 			taskToken,
 			session,
 			body,
-			stateMachine
+			stateMachine,
+			state
 		} = event;
 
 		try {
 
 			const data = await this.getFullData(body);
 
-			const dispatcher = new Dispatcher(LambdaFunction, session, data, taskToken);
+			const dispatcher = new Dispatcher(LambdaFunction, session, data, taskToken, state);
 
 			dispatcher.prepare();
 
@@ -66,7 +67,8 @@ module.exports = class Handler {
 
 			return {
 				...result,
-				stateMachine
+				stateMachine,
+				...state && { state }
 			};
 
 		} catch(error) {

diff --git a/lib/helpers/dispatcher.js b/lib/helpers/dispatcher.js
@@ -15,8 +15,9 @@ module.exports = class Dispatcher {
 	 * @param {ApiSession} session
 	 * @param {*} data
 	 * @param {string?} taskToken The task async token in case a function is invoked as an async task of a state machine
+	 * @param {*} state
 	 */
-	constructor(LambdaFunction, session, data, taskToken) {
+	constructor(LambdaFunction, session, data, taskToken, state) {
 
 		this.validateLambdaFunction(LambdaFunction);
 
@@ -28,6 +29,7 @@ module.exports = class Dispatcher {
 		this.session = session;
 		this.data = data;
 		this.taskToken = taskToken;
+		this.state = state;
 	}
 
 	/**
@@ -91,6 +93,7 @@ module.exports = class Dispatcher {
 		this.setSession(this.session);
 		this.setData(this.data);
 		this.setToken(this.taskToken);
+		this.setState(this.state);
 	}
 
 	setSession(session) {
@@ -105,6 +108,10 @@ module.exports = class Dispatcher {
 		this.lambdaFunction.taskToken = taskToken;
 	}
 
+	setState(state) {
+		this.lambdaFunction.state = state;
+	}
+
 	/**
 	 * Execute Lambda Function validate method
 	 */

diff --git a/lib/lambda-bases/lambda.js b/lib/lambda-bases/lambda.js
@@ -54,6 +54,21 @@ module.exports = class Lambda {
 		return this._taskToken;
 	}
 
+	/**
+	 * @param {*} state The invocation state. Only present for async invoked tasks of state machines
+	 */
+	set state(state) {
+		/** @private */
+		this._state = state;
+	}
+
+	/**
+	 *  @returns {*} The invocation state. Only present for async invoked tasks of state machines
+	*/
+	get state() {
+		return this._state;
+	}
+
 	/**
 	 * @returns {boolean} Whether the invocation must have a client associated or not
 	 */

diff --git a/lib/parallel-handler.js b/lib/parallel-handler.js
@@ -10,14 +10,17 @@ module.exports = class ParallelHandler extends Handler {
 	 * @returns {import('./handler').AWSLambdaEvent}
 	 */
 	static prepareEvent(event) {
-		return event.reduce((preparedEvent, { body, session, stateMachine }) => {
+		return event.reduce((preparedEvent, { body, session, stateMachine, state }) => {
 
 			if(session && !preparedEvent.session)
 				preparedEvent.session = session;
 
 			if(stateMachine && !preparedEvent.stateMachine)
 				preparedEvent.stateMachine = stateMachine;
 
+			if(state && !preparedEvent.state)
+				preparedEvent.state = state;
+
 			preparedEvent.body.push(body);
 
 			return preparedEvent;

diff --git a/tests/handler.js b/tests/handler.js
@@ -58,6 +58,17 @@ describe('Handler', () => {
 
 	const session = { clientCode: 'defaultClient' };
 
+	const stateMachine = {
+		Id: 'id-state-machine',
+		Name: 'state-machine-test'
+	};
+
+	const state = {
+		EnteredTime: '2019-03-26T20:14:13.192Z',
+		Name: 'Test',
+		RetryCount: 3
+	};
+
 	context('When Invalid Args is passed and default handling Validate Errors', () => {
 
 		it('Should return an error message if no Lambda Function is passed', async () => {
@@ -362,11 +373,6 @@ describe('Handler', () => {
 				contentS3Path
 			});
 
-			const stateMachine = {
-				id: 'id-state-machine',
-				name: 'state-machine-test'
-			};
-
 			const apiSession = new ApiSession({ ...session });
 			class LambdaFunctionExample {
 
@@ -380,8 +386,8 @@ describe('Handler', () => {
 
 			assert.deepStrictEqual(await Handler.handle(
 				LambdaFunctionExample,
-				{ session, body: { contentS3Path }, stateMachine }
-			), { session: apiSession, body: { contentS3Path }, stateMachine });
+				{ session, body: { contentS3Path }, stateMachine, state }
+			), { session: apiSession, body: { contentS3Path }, stateMachine, state });
 
 			sinon.assert.calledOnceWithExactly(Lambda.getBodyFromS3, contentS3Path);
 			sinon.assert.calledOnceWithExactly(Lambda.bodyToS3Path, 'step-function-payloads', body, []);
@@ -408,11 +414,6 @@ describe('Handler', () => {
 				contentS3Path
 			});
 
-			const stateMachine = {
-				id: 'id-state-machine',
-				name: 'state-machine-test'
-			};
-
 			const apiSession = new ApiSession({ ...session });
 			class LambdaFunctionExample {
 
@@ -426,8 +427,13 @@ describe('Handler', () => {
 
 			assert.deepStrictEqual(await Handler.handle(
 				LambdaFunctionExample,
-				{ session, body: { contentS3Path, error: { Error: 'Lambda.Unknown' } }, stateMachine }
-			), { session: apiSession, body: { contentS3Path }, stateMachine });
+				{
+					session,
+					body: { contentS3Path, error: { Error: 'Lambda.Unknown' } },
+					stateMachine,
+					state
+				}
+			), { session: apiSession, body: { contentS3Path }, stateMachine, state });
 
 			sinon.assert.calledOnceWithExactly(Lambda.getBodyFromS3, contentS3Path);
 			sinon.assert.calledOnceWithExactly(Lambda.bodyToS3Path, 'step-function-payloads', {
@@ -438,10 +444,6 @@ describe('Handler', () => {
 
 		it('Should return the same value when the payload has no session and body and the lambda is executed as a step function', async () => {
 
-			const stateMachine = {
-				id: 'id-state-machine',
-				name: 'state-machine-test'
-			};
 			class LambdaFunctionExample {
 
 				process() {
@@ -451,16 +453,12 @@ describe('Handler', () => {
 
 			assert.deepStrictEqual(await Handler.handle(
 				LambdaFunctionExample,
-				{ stateMachine }
+				{ stateMachine, state }
 			), {});
 		});
 
 		it('Should use a default value when the payload has no body and the lambda is executed as a step function', async () => {
 
-			const stateMachine = {
-				id: 'id-state-machine',
-				name: 'state-machine-test'
-			};
 			class LambdaFunctionExample {
 
 				process() {
@@ -472,16 +470,12 @@ describe('Handler', () => {
 
 			assert.deepStrictEqual(await Handler.handle(
 				LambdaFunctionExample,
-				{ stateMachine }
-			), { session: 'session', body: null, stateMachine });
+				{ stateMachine, state }
+			), { session: 'session', body: null, stateMachine, state });
 		});
 
 		it('Should use a default value when the payload has no sessions and the lambda is executed as a step function', async () => {
 
-			const stateMachine = {
-				id: 'id-state-machine',
-				name: 'state-machine-test'
-			};
 			class LambdaFunctionExample {
 
 				process() {
@@ -493,8 +487,8 @@ describe('Handler', () => {
 
 			assert.deepStrictEqual(await Handler.handle(
 				LambdaFunctionExample,
-				{ stateMachine }
-			), { session: null, body: {}, stateMachine });
+				{ stateMachine, state }
+			), { session: null, body: {}, stateMachine, state });
 		});
 	});
 });
diff --git a/tests/lambda-bases/lambda.js b/tests/lambda-bases/lambda.js
@@ -47,6 +47,21 @@ describe('Lambda bases', () => {
 		assert.deepStrictEqual(lambda.taskToken, invocationData);
 	});
 
+	it('Should have working getter and setter for state', () => {
+
+		const invocationData = {
+			EnteredTime: '2019-03-26T20:14:13.192Z',
+			Name: 'Test',
+			RetryCount: 3
+		};
+
+		const lambda = new Lambda();
+
+		lambda.state = invocationData;
+
+		assert.deepStrictEqual(lambda.state, invocationData);
+	});
+
 	it('Should have validation getter mustHaveClient set as false by default', () => {
 		const lambda = new Lambda();
 		assert.deepStrictEqual(lambda.mustHaveClient, false);

diff --git a/tests/parallel-handler.js b/tests/parallel-handler.js
@@ -10,6 +10,17 @@ describe('ParallelHandler', () => {
 
 	const session = { clientCode: 'defaultClient' };
 
+	const stateMachine = {
+		Id: 'id-state-machine',
+		Name: 'state-machine-test'
+	};
+
+	const state = {
+		EnteredTime: '2019-03-26T20:14:13.192Z',
+		Name: 'Test',
+		RetryCount: 3
+	};
+
 	context('When valid Args is passed', () => {
 
 		it('Should set correct session and data', async () => {
@@ -26,14 +37,9 @@ describe('ParallelHandler', () => {
 				pets: ['Dogs', 'Frogs']
 			};
 
-			const stateMachine = {
-				id: 'id-state-machine',
-				name: 'state-machine-test'
-			};
-
 			const event = [
-				{ body: body1, session, stateMachine },
-				{ body: body2, session, stateMachine }
+				{ body: body1, session, stateMachine, state },
+				{ body: body2, session, stateMachine, state }
 			];
 
 			const apiSession = new ApiSession({ ...session });
@@ -51,7 +57,8 @@ describe('ParallelHandler', () => {
 			assert.deepStrictEqual(await ParallelHandler.handle(LambdaFunctionExample, event), {
 				session: apiSession,
 				body: [body1, body2],
-				stateMachine
+				stateMachine,
+				state
 			});
 		});
 
@@ -99,16 +106,11 @@ describe('ParallelHandler', () => {
 				.withArgs(bodyLong2.contentS3Path)
 				.resolves(body4);
 
-			const stateMachine = {
-				id: 'id-state-machine',
-				name: 'state-machine-test'
-			};
-
 			const event = [
-				{ body: body1, session, stateMachine },
-				{ body: body2, session, stateMachine },
-				{ body: bodyLong1, session, stateMachine },
-				{ body: bodyLong2, session, stateMachine }
+				{ body: body1, session, stateMachine, state },
+				{ body: body2, session, stateMachine, state },
+				{ body: bodyLong1, session, stateMachine, state },
+				{ body: bodyLong2, session, stateMachine, state }
 			];
 
 			const apiSession = new ApiSession({ ...session });
@@ -127,7 +129,8 @@ describe('ParallelHandler', () => {
 			assert.deepStrictEqual(await ParallelHandler.handle(LambdaFunctionExample, event), {
 				session: apiSession,
 				body: [body1, body2, body3, body4],
-				stateMachine
+				stateMachine,
+				state
 			});
 		});
 	});