Inject context for xfn->Lambda when no Payload exists

DataDog · Sep 18, 2024 · f7222ff · f7222ff
1 parent f6d78ec
commit f7222ff
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
   }