Inject context for xfn->Lambda when no Payload exists

DataDog · Sep 18, 2024 · 2698227 · 2698227
1 parent f6d78ec
commit 2698227
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -53,7 +53,7 @@ describe('stepfunctions command helpers tests', () => {
  End: true,
  }
  expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeTruthy()
- expect(step.Parameters?.['Payload.$']).toEqual('States.JsonMerge($$, $, false)')
+ expect(step.Parameters?.['Payload.$']).toEqual(`$$['Execution', 'State', 'StateMachine']`)
  })
 
  test('default payload field of $', () => {

diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
@@ -110,10 +110,6 @@ export const injectContextIntoTasks = async (
  }
 }
 
-export const addTraceContextToLambdaParameters = (Parameters: ParametersType): void => {
- Parameters[`Payload.$`] = 'States.JsonMerge($$, $, false)'
-}
-
 export const addTraceContextToStepFunctionParameters = ({Parameters}: StepType): void => {
  if (Parameters) {
  if (!Parameters.Input) {
@@ -236,14 +232,14 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
 
  // payload field not set
  if (!step.Parameters.hasOwnProperty('Payload.$') && !step.Parameters.hasOwnProperty('Payload')) {
- addTraceContextToLambdaParameters(step.Parameters)
+ step.Parameters[`Payload.$`] = `$$['Execution', 'State', 'StateMachine']`
 
  return true
  }
 
  // default payload
  if (step.Parameters['Payload.$'] === '$') {
- addTraceContextToLambdaParameters(step.Parameters)
+ step.Parameters[`Payload.$`] = 'States.JsonMerge($$, $, false)'
 
  return true
  }