Refactor xfn Lambda injection code my merging two functions

DataDog · Sep 16, 2024 · 3566b41 · 3566b41
1 parent c48ed11
commit 3566b41
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 60 deletions.
diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -9,7 +9,7 @@ import {
   getStepFunctionLogGroupArn,
   parseArn,
   buildLogAccessPolicyName,
-  shouldUpdateStepForTracesMerging,
+  injectContextForLambdaFunctions,
   StepType,
   injectContextForStepFunctions,
   shouldUpdateStepForStepFunctionContextInjection,
@@ -29,7 +29,7 @@ const createMockContext = (): BaseContext => {
 
 describe('stepfunctions command helpers tests', () => {
   const context = createMockContext()
-  describe('shouldUpdateStepForTracesMerging test', () => {
+  describe('injectContextForLambdaFunctions test', () => {
     test('already has JsonMerge added to payload field', () => {
       const step: StepType = {
         Type: 'Task',
@@ -40,7 +40,7 @@ describe('stepfunctions command helpers tests', () => {
         },
         End: true,
       }
-      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
     })
 
     test('no payload field', () => {
@@ -52,7 +52,8 @@ describe('stepfunctions command helpers tests', () => {
         },
         End: true,
       }
-      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeTruthy()
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeTruthy()
+      expect(step.Parameters?.['Payload.$']).toEqual('States.JsonMerge($$, $, false)')
     })
 
     test('default payload field of $', () => {
@@ -65,7 +66,8 @@ describe('stepfunctions command helpers tests', () => {
         },
         End: true,
       }
-      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeTruthy()
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeTruthy()
+      expect(step.Parameters?.['Payload.$']).toEqual('States.JsonMerge($$, $, false)')
     })
 
     test('custom payload field not using JsonPath expression', () => {
@@ -78,7 +80,7 @@ describe('stepfunctions command helpers tests', () => {
         },
         End: true,
       }
-      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
     })
 
     test('custom payload field using JsonPath expression', () => {
@@ -91,7 +93,7 @@ describe('stepfunctions command helpers tests', () => {
         },
         End: true,
       }
-      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
+      expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
     })
 
     test('none-lambda step should not be updated', () => {
@@ -103,7 +105,7 @@ describe('stepfunctions command helpers tests', () => {
         },
         End: true,
       }
-      expect(shouldUpdateStepForTracesMerging(step, context, 'DynamoDB Update')).toBeFalsy()
+      expect(injectContextForLambdaFunctions(step, context, 'DynamoDB Update')).toBeFalsy()
     })
 
     test('legacy lambda api should not be updated', () => {
@@ -112,7 +114,7 @@ describe('stepfunctions command helpers tests', () => {
         Resource: 'arn:aws:lambda:sa-east-1:601427271234:function:hello-function',
         End: true,
       }
-      expect(shouldUpdateStepForTracesMerging(step, context, 'Legacy Lambda')).toBeFalsy()
+      expect(injectContextForLambdaFunctions(step, context, 'Legacy Lambda')).toBeFalsy()
     })
   })
 

diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
@@ -110,10 +110,8 @@ export const injectContextIntoTasks = async (
   }
 }
 
-export const addTraceContextToLambdaParameters = ({Parameters}: StepType): void => {
-  if (Parameters) {
-    Parameters[`Payload.$`] = 'States.JsonMerge($$, $, false)'
-  }
+export const addTraceContextToLambdaParameters = (Parameters: ParametersType): void => {
+  Parameters[`Payload.$`] = 'States.JsonMerge($$, $, false)'
 }
 
 export const addTraceContextToStepFunctionParameters = ({Parameters}: StepType): void => {
@@ -125,47 +123,6 @@ export const addTraceContextToStepFunctionParameters = ({Parameters}: StepType):
   }
 }
 
-export const shouldUpdateStepForTracesMerging = (step: StepType, context: BaseContext, stepName: string): boolean => {
-  // not default lambda api
-  if (step.Resource !== 'arn:aws:states:::lambda:invoke') {
-    return false
-  }
-
-  if (!step.Parameters) {
-    context.stdout
-      .write(`[Warn] Step ${stepName} does not have a Parameters field. Step Functions Context Object injection \
-skipped. Your Step Functions trace will not be merged with downstream Lambda traces. To manually merge these traces, \
-check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
-
-    return false
-  }
-
-  // payload field not set
-  if (!step.Parameters.hasOwnProperty('Payload.$') && !step.Parameters.hasOwnProperty('Payload')) {
-    return true
-  }
-
-  // default payload
-  if (step.Parameters['Payload.$'] === '$') {
-    return true
-  }
-
-  // context injection is already set up
-  if (step.Parameters['Payload.$'] === 'States.JsonMerge($$, $, false)') {
-    context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
-
-    return false
-  }
-
-  // custom payload
-  context.stdout
-    .write(`[Warn] Step ${stepName} has a custom Payload field. Step Functions Context Object injection skipped. \
-Your Step Functions trace will not be merged with downstream Lambda traces. To manually merge these traces, \
-check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
-
-  return false
-}
-
 // Truth table
 // Input                    | Expected
 // -------------------------|---------
@@ -252,19 +209,58 @@ export type ParametersType = {
   }
 }
 
-const injectContextForLambdaFunctions = (step: StepType, context: BaseContext, stepName: string): boolean => {
-  if (shouldUpdateStepForTracesMerging(step, context, stepName)) {
-    addTraceContextToLambdaParameters(step)
-
-    return true
-  } else if (step.Resource?.startsWith('arn:aws:lambda')) {
+export const injectContextForLambdaFunctions = (step: StepType, context: BaseContext, stepName: string): boolean => {
+  if (step.Resource?.startsWith('arn:aws:lambda')) {
     context.stdout.write(
       `[Warn] Step ${stepName} may be using the basic legacy integration, which does not support merging lambda trace(s) with Step Functions trace.
           To merge lambda trace(s) with Step Functions trace, please consider using the latest integration.
           More details can be found on https://docs.aws.amazon.com/step-functions/latest/dg/connect-lambda.html \n`
     )
+
+    return false
+  }
+
+  // not default lambda api
+  if (step.Resource !== 'arn:aws:states:::lambda:invoke') {
+    return false
+  }
+
+  if (!step.Parameters) {
+    context.stdout
+      .write(`[Warn] Step ${stepName} does not have a Parameters field. Step Functions Context Object injection \
+skipped. Your Step Functions trace will not be merged with downstream Lambda traces. To manually merge these traces, \
+check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
+
+    return false
+  }
+
+  // payload field not set
+  if (!step.Parameters.hasOwnProperty('Payload.$') && !step.Parameters.hasOwnProperty('Payload')) {
+    addTraceContextToLambdaParameters(step.Parameters)
+
+    return true
   }
 
+  // default payload
+  if (step.Parameters['Payload.$'] === '$') {
+    addTraceContextToLambdaParameters(step.Parameters)
+
+    return true
+  }
+
+  // context injection is already set up
+  if (step.Parameters['Payload.$'] === 'States.JsonMerge($$, $, false)') {
+    context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
+
+    return false
+  }
+
+  // custom payload
+  context.stdout
+    .write(`[Warn] Step ${stepName} has a custom Payload field. Step Functions Context Object injection skipped. \
+Your Step Functions trace will not be merged with downstream Lambda traces. To manually merge these traces, \
+check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
+
   return false
 }