fix(stepfunctions): the catch field in CustomState is not rendered

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.ts

@@ -23,11 +23,6 @@ const stateJson = {
     },
   },
   ResultPath: null,
-  Retry: [{
-    ErrorEquals: [sfn.Errors.PERMISSIONS],
-    IntervalSeconds: 20,
-    MaxAttempts: 2,
-  }],
 };
 
 const failure = new sfn.Fail(stack, 'failed', {
@@ -39,18 +34,35 @@ const custom = new sfn.CustomState(stack, 'my custom task', {
   stateJson,
 });
 
-const customWithInlineRetry = new sfn.CustomState(stack, 'my custom task with inline Retriers', {
-  stateJson,
-});
-
 custom.addCatch(failure);
 custom.addRetry({
   errors: [sfn.Errors.TIMEOUT],
   interval: cdk.Duration.seconds(10),
   maxAttempts: 5,
 });
 
-const chain = sfn.Chain.start(custom).next(customWithInlineRetry).next(finalStatus);
+const customWithInlineRetry = new sfn.CustomState(stack, 'my custom task with inline Retriers', {
+  stateJson: {
+    ...stateJson,
+    Retry: [{
+      ErrorEquals: [sfn.Errors.PERMISSIONS],
+      IntervalSeconds: 20,
+      MaxAttempts: 2,
+    }]
+  }
+});
+
+const customWithInlineCatch = new sfn.CustomState(stack, 'my custom task with inline Catchers', {
+  stateJson: {
+    ...stateJson,
+    Catch: [{
+      ErrorEquals: [sfn.Errors.PERMISSIONS],
+      Next: 'failed'
+    }]
+  }
+});
+
+const chain = sfn.Chain.start(custom).next(customWithInlineRetry).next(customWithInlineCatch).next(finalStatus);
 
 const sm = new sfn.StateMachine(stack, 'StateMachine', {
   definition: chain,

packages/aws-cdk-lib/aws-stepfunctions/lib/states/custom-state.ts

@@ -74,11 +74,16 @@ export class CustomState extends State implements IChainable, INextable {
       ...this.renderRetryCatch(),
     };
 
-    // merge the Retry field defined in the stateJson into the state
+    // Retriers and Catchers can be specified directly in the stateJson or indirectly to the construct with addRetry() and addCatch().
+    // renderRetryCatch() only renders the indirectly supplied Retriers and Catchers, so we need to manually merge in those directly in the stateJson
     if (Array.isArray(this.stateJson.Retry)) {
       state.Retry = Array.isArray(state.Retry) ? [...state.Retry, ...this.stateJson.Retry] : [...this.stateJson.Retry];
     }
 
+    if (Array.isArray(this.stateJson.Catch)) {
+      state.Catch = Array.isArray(state.Catch) ? [...state.Catch, ...this.stateJson.Catch] : [...this.stateJson.Catch];
+    }
+
     return state;
   }
 }

packages/aws-cdk-lib/aws-stepfunctions/test/custom-state.test.ts

@@ -1,8 +1,9 @@
 import { render } from './private/render-util';
 import * as cdk from '../../core';
 import * as sfn from '../lib';
+import { Errors } from '../lib/types';
 
-describe('Custom State', () => {
+describe.only('Custom State', () => {
   let stack: cdk.Stack;
   let stateJson: any;
 
@@ -309,4 +310,145 @@ describe('Custom State', () => {
       },
     );
   });
+
+  test('expect retry to merge when specifying strategy inline and through construct', () => {
+    // GIVEN
+    const custom = new sfn.CustomState(stack, 'Custom', {
+      stateJson: {
+        ...stateJson,
+        Retry: [{
+          ErrorEquals: ['States.TaskFailed'],
+        }],
+      },
+    }).addRetry({ errors: [Errors.TIMEOUT] });
+    const chain = sfn.Chain.start(custom);
+
+    // THEN
+    expect(render(stack, chain)).toStrictEqual(
+      {
+        StartAt: 'Custom',
+        States: {
+          Custom: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::dynamodb:putItem',
+            Parameters: {
+              TableName: 'MyTable',
+              Item: {
+                id: {
+                  S: 'MyEntry',
+                },
+              },
+            },
+            ResultPath: null,
+            Retry: [
+              {
+                ErrorEquals: ['States.Timeout'],
+              },
+              {
+                ErrorEquals: ['States.TaskFailed'],
+              },
+            ],
+            End: true,
+          },
+        },
+      },
+    );
+  });
+
+  test('expect catch to not fail when specifying strategy inline', () => {
+    // GIVEN
+    const custom = new sfn.CustomState(stack, 'Custom', {
+      stateJson: {
+        ...stateJson,
+        Catch: [{
+          ErrorEquals: ['States.TaskFailed'],
+          Next: 'Failed',
+        }],
+      },
+    });
+    const chain = sfn.Chain.start(custom);
+
+    // THEN
+    expect(render(stack, chain)).toStrictEqual(
+      {
+        StartAt: 'Custom',
+        States: {
+          Custom: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::dynamodb:putItem',
+            Parameters: {
+              TableName: 'MyTable',
+              Item: {
+                id: {
+                  S: 'MyEntry',
+                },
+              },
+            },
+            ResultPath: null,
+            Catch: [{
+              ErrorEquals: ['States.TaskFailed'],
+              Next: 'Failed',
+            }],
+            End: true,
+          },
+        },
+      },
+    );
+  });
+
+  test('expect catch to merge when specifying strategy inline and through construct', () => {
+    // GIVEN
+    const failure = new sfn.Fail(stack, 'Failed', {
+      error: 'DidNotWork',
+      cause: 'We got stuck',
+    });
+
+    const custom = new sfn.CustomState(stack, 'Custom', {
+      stateJson: {
+        ...stateJson,
+        Catch: [{
+          ErrorEquals: ['States.TaskFailed'],
+          Next: 'Failed',
+        }],
+      },
+    }).addCatch(failure, { errors: [Errors.TIMEOUT] });
+    const chain = sfn.Chain.start(custom);
+
+    // THEN
+    expect(render(stack, chain)).toStrictEqual(
+      {
+        StartAt: 'Custom',
+        States: {
+          Custom: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::dynamodb:putItem',
+            Parameters: {
+              TableName: 'MyTable',
+              Item: {
+                id: {
+                  S: 'MyEntry',
+                },
+              },
+            },
+            ResultPath: null,
+            Catch: [
+              {
+                ErrorEquals: ['States.Timeout'],
+                Next: 'Failed',
+              }, {
+                ErrorEquals: ['States.TaskFailed'],
+                Next: 'Failed',
+              },
+            ],
+            End: true,
+          },
+          Failed: {
+            Type: 'Fail',
+            Error: 'DidNotWork',
+            Cause: 'We got stuck',
+          },
+        },
+      },
+    );
+  });
 });