From 2f801513b8ec312277c57640f23ed5cb666ef8ff Mon Sep 17 00:00:00 2001
From: Michael Sambol <sambol.michael@gmail.com>
Date: Tue, 5 Sep 2023 13:34:32 -0500
Subject: [PATCH] Add feedback from Kaizen

---
 .../aws-stepfunctions/test/integ.state-machine.ts   |  4 +++-
 packages/aws-cdk-lib/aws-stepfunctions/README.md    |  4 +++-
 .../aws-stepfunctions/lib/states/choice.ts          |  6 +++---
 .../aws-stepfunctions/lib/states/state.ts           | 13 +++++++++----
 .../aws-stepfunctions/test/state-machine.test.ts    |  4 +++-
 5 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.ts
index 4d69c528eafc6..5ff6711106589 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.ts
@@ -24,7 +24,9 @@ const choice = new sfn.Choice(stack, 'choice', {
 
 const success = new sfn.Succeed(stack, 'success');
 
-choice.when(sfn.Condition.isPresent('$.success'), success, 'this is a comment for the when condition');
+choice.when(sfn.Condition.isPresent('$.success'), success, {
+  comment: 'this is a comment for the when condition',
+});
 choice.when(sfn.Condition.isPresent('$.noComment'), shortWait);
 choice.otherwise(success);
 wait.next(choice);
diff --git a/packages/aws-cdk-lib/aws-stepfunctions/README.md b/packages/aws-cdk-lib/aws-stepfunctions/README.md
index 51a1488815b09..b0af8ce337f37 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions/README.md
+++ b/packages/aws-cdk-lib/aws-stepfunctions/README.md
@@ -359,7 +359,9 @@ const choice = new sfn.Choice(this, 'What color is it?', {
 });
 const handleBlueItem = new sfn.Pass(this, 'HandleBlueItem');
 const handleOtherItemColor = new sfn.Pass(this, 'HanldeOtherItemColor');
-choice.when(sfn.Condition.stringEquals('$.color', 'BLUE'), handleBlueItem, 'blue item comment');
+choice.when(sfn.Condition.stringEquals('$.color', 'BLUE'), handleBlueItem, {
+  comment: 'blue item comment',
+});
 choice.otherwise(handleOtherItemColor);
 ```
 
diff --git a/packages/aws-cdk-lib/aws-stepfunctions/lib/states/choice.ts b/packages/aws-cdk-lib/aws-stepfunctions/lib/states/choice.ts
index f02f5c19e9e7f..cb876a239fae0 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions/lib/states/choice.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions/lib/states/choice.ts
@@ -1,6 +1,6 @@
 import { Construct } from 'constructs';
 import { StateType } from './private/state-type';
-import { State } from './state';
+import { ChoiceTransitionOptions, State } from './state';
 import { Chain } from '../chain';
 import { Condition } from '../condition';
 import { IChainable, INextable } from '../types';
@@ -53,8 +53,8 @@ export class Choice extends State {
   /**
    * If the given condition matches, continue execution with the given state
    */
-  public when(condition: Condition, next: IChainable, comment?: string): Choice {
-    super.addChoice(condition, next.startState, comment);
+  public when(condition: Condition, next: IChainable, options: ChoiceTransitionOptions = {}): Choice {
+    super.addChoice(condition, next.startState, options);
     return this;
   }
 
diff --git a/packages/aws-cdk-lib/aws-stepfunctions/lib/states/state.ts b/packages/aws-cdk-lib/aws-stepfunctions/lib/states/state.ts
index b9906e5328ea7..29eeba5c8a5ec 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions/lib/states/state.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions/lib/states/state.ts
@@ -318,8 +318,8 @@ export abstract class State extends Construct implements IChainable {
   /**
    * Add a choice branch to this state
    */
-  protected addChoice(condition: Condition, next: State, comment?: string) {
-    this.choices.push({ condition, next, comment });
+  protected addChoice(condition: Condition, next: State, options: ChoiceTransitionOptions) {
+    this.choices.push({ condition, next, ...options });
     next.startState.addIncoming(this);
     if (this.containingGraph) {
       next.startState.bindToGraph(this.containingGraph);
@@ -479,7 +479,7 @@ export interface FindStateOptions {
 /**
  * A Choice Transition
  */
-interface ChoiceTransition {
+interface ChoiceTransition extends ChoiceTransitionOptions {
   /**
    * State to transition to
    */
@@ -489,9 +489,14 @@ interface ChoiceTransition {
    * Condition for this transition
    */
   condition: Condition;
+}
 
+/**
+ * Options for Choice Transition
+ */
+export interface ChoiceTransitionOptions {
   /**
-   * An optional description for this choice transition
+   * An optional description for the choice transition
    *
    * @default No comment
    */
diff --git a/packages/aws-cdk-lib/aws-stepfunctions/test/state-machine.test.ts b/packages/aws-cdk-lib/aws-stepfunctions/test/state-machine.test.ts
index 15beb79986ded..df71c7d4fa392 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions/test/state-machine.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions/test/state-machine.test.ts
@@ -631,7 +631,9 @@ describe('State Machine', () => {
       comment: 'nebraska',
     });
     const success = new sfn.Succeed(stack, 'success');
-    choice.when(sfn.Condition.isPresent('$.success'), success, 'london');
+    choice.when(sfn.Condition.isPresent('$.success'), success, {
+      comment: 'london',
+    });
     choice.otherwise(success);
 
     // WHEN