Refactor sfn->sfn context injection code my merging two functions

src/commands/stepfunctions/__tests__/helpers.test.ts

@@ -12,7 +12,6 @@ import {
  injectContextForLambdaFunctions,
  StepType,
  injectContextForStepFunctions,
- shouldUpdateStepForStepFunctionContextInjection,
  PayloadObject,
 } from '../helpers'
 
@@ -277,8 +276,8 @@ describe('stepfunctions command helpers tests', () => {
  })
  })
 
- describe('shouldUpdateStepForStepFunctionContextInjection', () => {
- test('is true for an empty object', () => {
+ describe('inject context for StepFunctions', () => {
+ test('injects context into a step function invocation', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
@@ -288,71 +287,61 @@ describe('stepfunctions command helpers tests', () => {
  },
  End: true,
  }
- expect(
- shouldUpdateStepForStepFunctionContextInjection(step, context, 'Step Functions StartExecution')
- ).toBeTruthy()
+
+ const changed = injectContextForStepFunctions(step, context, 'Step Functions StartExecution')
+ expect(changed).toBeTruthy()
+ expect(step.Parameters?.Input).toEqual({'CONTEXT.$': 'States.JsonMerge($$, $, false)'})
  })
 
- test('is true for undefined', () => {
+ test('is true for an empty object', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
  Parameters: {
  StateMachineArn: 'arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine',
+ Input: {},
  },
  End: true,
  }
- expect(
- shouldUpdateStepForStepFunctionContextInjection(step, context, 'Step Functions StartExecution')
- ).toBeTruthy()
+ expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeTruthy()
  })
 
- test('is false when Input is an object that contains a CONTEXT key using JSONPath expression', () => {
+ test('is true for undefined', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
  Parameters: {
  StateMachineArn: 'arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine',
- Input: {'CONTEXT.$': 'blah'},
  },
  End: true,
  }
- expect(
- shouldUpdateStepForStepFunctionContextInjection(step, context, 'Step Functions StartExecution')
- ).toBeFalsy()
+ expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeTruthy()
  })
 
- test('is false when Input is an object that contains a CONTEXT key not using JSONPath expression', () => {
+ test('is false when Input is an object that contains a CONTEXT key using JSONPath expression', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
  Parameters: {
  StateMachineArn: 'arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine',
- Input: {CONTEXT: 'blah'},
+ Input: {'CONTEXT.$': 'blah'},
  },
  End: true,
  }
- expect(
- shouldUpdateStepForStepFunctionContextInjection(step, context, 'Step Functions StartExecution')
- ).toBeFalsy()
+ expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeFalsy()
  })
- })
 
- describe('inject context for StepFunctions', () => {
- test('injects context into a step function invocation', () => {
+ test('is false when Input is an object that contains a CONTEXT key not using JSONPath expression', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
  Parameters: {
  StateMachineArn: 'arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine',
- Input: {},
+ Input: {CONTEXT: 'blah'},
  },
  End: true,
  }
-
- const changed = injectContextForStepFunctions(step, context, 'Step Functions StartExecution')
- expect(changed).toBeTruthy()
- expect(step.Parameters?.Input).toEqual({'CONTEXT.$': 'States.JsonMerge($$, $, false)'})
+ expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeFalsy()
  })
  })
 

src/commands/stepfunctions/helpers.ts

@@ -119,65 +119,6 @@ export const addTraceContextToStepFunctionParameters = ({Parameters}: StepType):
  }
 }
 
-// Truth table
-// Input | Expected
-// -------------------------|---------
-// Empty object | true
-// undefined | true
-// not object | false
-// object without CONTEXT.$ | true
-// object with CONTEXT.$ | false
-export const shouldUpdateStepForStepFunctionContextInjection = (
- step: StepType,
- context: BaseContext,
- stepName: string
-): boolean => {
- // not default lambda api
- if (!step.Resource?.startsWith('arn:aws:states:::states:startExecution')) {
- 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 Step Function traces. To manually merge these \
-traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
-
- return false
- }
-
- if (!step.Parameters.Input) {
- return true
- }
-
- if (typeof step.Parameters.Input !== 'object') {
- context.stdout
- .write(`[Warn] Step ${stepName}'s Parameters.Input field is not a JSON object. Step Functions Context Object \
-injection skipped. Your Step Functions trace will not be merged with downstream Step Function traces. To manually \
-merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
-
- return false
- }
-
- if (!step.Parameters.Input['CONTEXT.$'] && !step.Parameters.Input['CONTEXT']) {
- return true
- }
-
- // context injection is already set up
- if (step.Parameters.Input['CONTEXT.$'] === 'States.JsonMerge($$, $, false)') {
- context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
-
- return false
- }
-
- context.stdout
- .write(`[Warn] Step ${stepName}'s Parameters.Input field has a custom CONTEXT field. Step Functions Context \
-Object injection skipped. Your Step Functions trace will not be merged with downstream Step Function traces. To \
-manually merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
-
- return false
-}
-
 export type StateMachineDefinitionType = {
  Comment?: string
  StartAt?: string
@@ -328,11 +269,52 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
 }
 
 export const injectContextForStepFunctions = (step: StepType, context: BaseContext, stepName: string): boolean => {
- if (shouldUpdateStepForStepFunctionContextInjection(step, context, stepName)) {
+ // not default lambda api
+ if (!step.Resource?.startsWith('arn:aws:states:::states:startExecution')) {
+ 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 Step Function traces. To manually merge these \
+traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
+
+ return false
+ }
+
+ if (!step.Parameters.Input) {
  addTraceContextToStepFunctionParameters(step)
 
  return true
  }
 
+ if (typeof step.Parameters.Input !== 'object') {
+ context.stdout
+ .write(`[Warn] Step ${stepName}'s Parameters.Input field is not a JSON object. Step Functions Context Object \
+injection skipped. Your Step Functions trace will not be merged with downstream Step Function traces. To manually \
+merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
+
+ return false
+ }
+
+ if (!step.Parameters.Input['CONTEXT.$'] && !step.Parameters.Input['CONTEXT']) {
+ addTraceContextToStepFunctionParameters(step)
+
+ return true
+ }
+
+ // context injection is already set up
+ if (step.Parameters.Input['CONTEXT.$'] === 'States.JsonMerge($$, $, false)') {
+ context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
+
+ return false
+ }
+
+ context.stdout
+ .write(`[Warn] Step ${stepName}'s Parameters.Input field has a custom CONTEXT field. Step Functions Context \
+Object injection skipped. Your Step Functions trace will not be merged with downstream Step Function traces. To \
+manually merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)
+
  return false
 }