diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts index c493a104b..8352387a8 100644 --- a/src/commands/stepfunctions/__tests__/helpers.test.ts +++ b/src/commands/stepfunctions/__tests__/helpers.test.ts @@ -43,6 +43,19 @@ describe('stepfunctions command helpers tests', () => { expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy() }) + test('Case 4.2: context injection already set up', () => { + const step: StepType = { + Type: 'Task', + Resource: 'arn:aws:states:::lambda:invoke', + Parameters: { + FunctionName: 'arn:aws:lambda:sa-east-1:425362991234:function:unit-test-lambda-function', + 'Payload.$': `$$['Execution', 'State', 'StateMachine']`, + }, + End: true, + } + expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy() + }) + test('Case 1: no Payload or Payload.$', () => { const step: StepType = { Type: 'Task', diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts index c28b8034e..f1e651ace 100644 --- a/src/commands/stepfunctions/helpers.ts +++ b/src/commands/stepfunctions/helpers.ts @@ -155,7 +155,8 @@ export type ParametersType = { // 2.3 | "Payload" object has no Execution, State or StateMachine | true // 3 | "Payload" is not object | false // 4.1 | "Payload.$": "$" (default payload) | true -// 4.2 | "Payload.$": "States.JsonMerge($$, $, false)" | false +// 4.2 | "Payload.$": "States.JsonMerge($$, $, false)" or | false +// | "Payload.$": "$$['Execution', 'State', 'StateMachine']" | // 4.3 | Custom "Payload.$" | false export const injectContextForLambdaFunctions = (step: StepType, context: BaseContext, stepName: string): boolean => { if (step.Resource?.startsWith('arn:aws:lambda')) { @@ -244,7 +245,10 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio } // Case 4.2: context injection is already set up using "Payload.$" - if (step.Parameters['Payload.$'] === 'States.JsonMerge($$, $, false)') { + if ( + step.Parameters['Payload.$'] === 'States.JsonMerge($$, $, false)' || + step.Parameters['Payload.$'] === `$$['Execution', 'State', 'StateMachine']` + ) { context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`) return false