Handle another case when sfn->Lambda context injection is already set…

… up (case 4.2)
DataDog · Sep 26, 2024 · 2cc8717 · 2cc8717
1 parent d36a07e
commit 2cc8717
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git 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
@@ -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