From 7afa665a9b041205824d24f915f74073c27b7639 Mon Sep 17 00:00:00 2001
From: Al Harris <91494052+alharris-at@users.noreply.github.com>
Date: Wed, 9 Feb 2022 14:37:41 -0800
Subject: [PATCH] chore: add control overrides for checkbox and switch (#390)

---
 .../studio-ui-codegen-react.test.ts.snap      | 105 ++++++++++++++++--
 .../__tests__/studio-ui-codegen-react.test.ts |   8 ++
 .../lib/react-component-renderer.ts           |   1 +
 .../react-component-with-children-renderer.ts |   1 +
 .../codegen-ui-react/lib/workflow/mutation.ts |  99 ++++++++++-------
 .../workflow/checkboxControlledElement.json   |  24 ++++
 .../workflow/switchControlledElement.json     |  24 ++++
 .../workflow/twoWayBindings.json              |   4 +-
 .../components/workflow/twoWayBindings.json   |   4 +-
 9 files changed, 220 insertions(+), 50 deletions(-)
 create mode 100644 packages/codegen-ui/example-schemas/workflow/checkboxControlledElement.json
 create mode 100644 packages/codegen-ui/example-schemas/workflow/switchControlledElement.json

diff --git a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react.test.ts.snap b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react.test.ts.snap
index ac3272ce3..68ae1eb87 100644
--- a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react.test.ts.snap
+++ b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react.test.ts.snap
@@ -5406,6 +5406,53 @@ export default function TwoWayBindings(
 }
 `;
 
+exports[`amplify render tests mutations supports a controlled checkbox primitive 1`] = `
+Object {
+  "componentText": "/* eslint-disable */
+import React from \\"react\\";
+import {
+  EscapeHatchProps,
+  getOverrideProps,
+  useStateMutationAction,
+} from \\"@aws-amplify/ui-react/internal\\";
+import { CheckboxField, Flex, FlexProps, Text } from \\"@aws-amplify/ui-react\\";
+
+export type CheckboxControlledElementProps = React.PropsWithChildren<
+  Partial<FlexProps> & {
+    overrides?: EscapeHatchProps | undefined | null;
+  }
+>;
+export default function CheckboxControlledElement(
+  props: CheckboxControlledElementProps
+): React.ReactElement {
+  const { overrides, ...rest } = props;
+  const [inputChecked, setInputChecked] = useStateMutationAction(undefined);
+  return (
+    /* @ts-ignore: TS2322 */
+    <Flex
+      {...rest}
+      {...getOverrideProps(overrides, \\"CheckboxControlledElement\\")}
+    >
+      <CheckboxField
+        checked={inputChecked}
+        onChange={(event: SyntheticEvent) => {
+          setInputChecked(event.target.checked);
+        }}
+        {...getOverrideProps(overrides, \\"Input\\")}
+      ></CheckboxField>
+      <Text
+        children={inputChecked}
+        {...getOverrideProps(overrides, \\"CheckboxFieldValue\\")}
+      ></Text>
+    </Flex>
+  );
+}
+",
+  "declaration": undefined,
+  "renderComponentToFilesystem": [Function],
+}
+`;
+
 exports[`amplify render tests mutations supports a controlled stepper primitive 1`] = `
 Object {
   "componentText": "/* eslint-disable */
@@ -5457,6 +5504,48 @@ export default function StepperControlledElement(
 }
 `;
 
+exports[`amplify render tests mutations supports a controlled switch primitive 1`] = `
+Object {
+  "componentText": "/* eslint-disable */
+import React from \\"react\\";
+import {
+  EscapeHatchProps,
+  getOverrideProps,
+  useStateMutationAction,
+} from \\"@aws-amplify/ui-react/internal\\";
+import { Flex, FlexProps, SwitchField, Text } from \\"@aws-amplify/ui-react\\";
+
+export type SwitchControlledElementProps = React.PropsWithChildren<
+  Partial<FlexProps> & {
+    overrides?: EscapeHatchProps | undefined | null;
+  }
+>;
+export default function SwitchControlledElement(
+  props: SwitchControlledElementProps
+): React.ReactElement {
+  const { overrides, ...rest } = props;
+  const [inputIsChecked, setInputIsChecked] = useStateMutationAction(undefined);
+  return (
+    /* @ts-ignore: TS2322 */
+    <Flex {...rest} {...getOverrideProps(overrides, \\"SwitchControlledElement\\")}>
+      <SwitchField
+        isChecked={inputIsChecked}
+        onChange={() => setInputIsChecked(!inputIsChecked)}
+        {...getOverrideProps(overrides, \\"Input\\")}
+      ></SwitchField>
+      <Text
+        children={inputIsChecked}
+        {...getOverrideProps(overrides, \\"SwitchFieldValue\\")}
+      ></Text>
+    </Flex>
+  );
+}
+",
+  "declaration": undefined,
+  "renderComponentToFilesystem": [Function],
+}
+`;
+
 exports[`amplify render tests mutations supports multiple actions pointing to the same value 1`] = `
 Object {
   "componentText": "/* eslint-disable */
@@ -5741,7 +5830,7 @@ export default function TwoWayBindings(
     useStateMutationAction(undefined);
   const [stepperFieldInputValue, setStepperFieldInputValue] =
     useStateMutationAction(undefined);
-  const [switchFieldInputValue, setSwitchFieldInputValue] =
+  const [switchFieldInputIsChecked, setSwitchFieldInputIsChecked] =
     useStateMutationAction(undefined);
   const [textFieldInputValue, setTextFieldInputValue] =
     useStateMutationAction(undefined);
@@ -5770,7 +5859,7 @@ export default function TwoWayBindings(
     setStepperFieldInputValue(9);
   };
   const setSwitchFieldValueClick = () => {
-    setSwitchFieldInputValue(true);
+    setSwitchFieldInputIsChecked(true);
   };
   const setTextFieldValueClick = () => {
     setTextFieldInputValue(\\"Hardcoded Value\\");
@@ -5802,7 +5891,7 @@ export default function TwoWayBindings(
           value=\\"yes\\"
           checked={checkboxFieldInputChecked}
           onChange={(event: SyntheticEvent) => {
-            setCheckboxFieldInputChecked(event.target.value);
+            setCheckboxFieldInputChecked(event.target.checked);
           }}
           {...getOverrideProps(overrides, \\"CheckboxFieldInput\\")}
         ></CheckboxField>
@@ -6079,14 +6168,14 @@ export default function TwoWayBindings(
           labelPosition=\\"start\\"
           isDisabled={false}
           isLabelHidden={false}
-          value={switchFieldInputValue}
-          onChange={(event: SyntheticEvent) => {
-            setSwitchFieldInputValue(event.target.value);
-          }}
+          isChecked={switchFieldInputIsChecked}
+          onChange={() =>
+            setSwitchFieldInputIsChecked(!switchFieldInputIsChecked)
+          }
           {...getOverrideProps(overrides, \\"SwitchFieldInput\\")}
         ></SwitchField>
         <Text
-          children={switchFieldInputValue}
+          children={switchFieldInputIsChecked}
           {...getOverrideProps(overrides, \\"SwitchFieldValue\\")}
         ></Text>
         <Button
diff --git a/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react.test.ts b/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react.test.ts
index 81ea650a0..e014c2e8b 100644
--- a/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react.test.ts
+++ b/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react.test.ts
@@ -362,6 +362,14 @@ describe('amplify render tests', () => {
       expect(generateWithAmplifyRenderer('workflow/stepperControlledElement')).toMatchSnapshot();
     });
 
+    it('supports a controlled checkbox primitive', () => {
+      expect(generateWithAmplifyRenderer('workflow/checkboxControlledElement')).toMatchSnapshot();
+    });
+
+    it('supports a controlled switch primitive', () => {
+      expect(generateWithAmplifyRenderer('workflow/switchControlledElement')).toMatchSnapshot();
+    });
+
     it('modifies text in component on input change', () => {
       expect(generateWithAmplifyRenderer('workflow/inputToTextChange')).toMatchSnapshot();
     });
diff --git a/packages/codegen-ui-react/lib/react-component-renderer.ts b/packages/codegen-ui-react/lib/react-component-renderer.ts
index 5ab731107..358a18aa1 100644
--- a/packages/codegen-ui-react/lib/react-component-renderer.ts
+++ b/packages/codegen-ui-react/lib/react-component-renderer.ts
@@ -96,6 +96,7 @@ export class ReactComponentRenderer<TPropIn> extends ComponentRendererBase<
         buildOpeningElementControlEvents(
           this.component.componentType,
           getSetStateName({ componentName: this.component.name || '', property: key }),
+          getStateName({ componentName: this.component.name || '', property: key }),
           'change',
         ),
       );
diff --git a/packages/codegen-ui-react/lib/react-component-with-children-renderer.ts b/packages/codegen-ui-react/lib/react-component-with-children-renderer.ts
index 2be3e76f2..7606a6a74 100644
--- a/packages/codegen-ui-react/lib/react-component-with-children-renderer.ts
+++ b/packages/codegen-ui-react/lib/react-component-with-children-renderer.ts
@@ -95,6 +95,7 @@ export class ReactComponentWithChildrenRenderer<TPropIn> extends ComponentWithCh
         buildOpeningElementControlEvents(
           this.component.componentType,
           getSetStateName({ componentName: this.component.name || '', property: key }),
+          getStateName({ componentName: this.component.name || '', property: key }),
           'change',
         ),
       );
diff --git a/packages/codegen-ui-react/lib/workflow/mutation.ts b/packages/codegen-ui-react/lib/workflow/mutation.ts
index 7e6cfee87..fc2aab45f 100644
--- a/packages/codegen-ui-react/lib/workflow/mutation.ts
+++ b/packages/codegen-ui-react/lib/workflow/mutation.ts
@@ -32,15 +32,17 @@ import {
   getSetStateName,
 } from '../react-component-render-helper';
 import { ImportCollection, ImportValue } from '../imports';
+import Primitive, { PrimitivesWithChangeEvent, PrimitiveLevelPropConfiguration } from '../primitive';
 import { mapGenericEventToReact } from './events';
 import { getChildPropMappingForComponentName } from './utils';
-import Primitive, { PrimitivesWithChangeEvent, PrimitiveLevelPropConfiguration } from '../primitive';
 
-type EventHandlerBuilder = (stateName: string) => JsxExpression;
+type EventHandlerBuilder = (setStateName: string, stateName: string) => JsxExpression;
 
 const genericEventToReactEventImplementationOverrides: PrimitiveLevelPropConfiguration<EventHandlerBuilder> = {
-  [Primitive.StepperField]: { [StudioGenericEvent.change]: numericValueCallbackGenerator },
-  [Primitive.SliderField]: { [StudioGenericEvent.change]: numericValueCallbackGenerator },
+  [Primitive.StepperField]: { [StudioGenericEvent.change]: numericValueCallback },
+  [Primitive.SliderField]: { [StudioGenericEvent.change]: numericValueCallback },
+  [Primitive.CheckboxField]: { [StudioGenericEvent.change]: buildSyntheticEventTargetCallback('checked') },
+  [Primitive.SwitchField]: { [StudioGenericEvent.change]: toggleBooleanStateCallback },
 };
 
 export function getComponentStateReferences(component: StudioComponent) {
@@ -107,7 +109,48 @@ export function getActionStateParameters(action: ActionStudioComponentEvent): St
   return [];
 }
 
-function syntheticEventTargetValueCallbackGenerator(stateName: string): JsxExpression {
+function buildSyntheticEventTargetCallback(targetFieldName: string): (setStateName: string) => JsxExpression {
+  return (setStateName: string) => {
+    return factory.createJsxExpression(
+      undefined,
+      factory.createArrowFunction(
+        undefined,
+        undefined,
+        [
+          factory.createParameterDeclaration(
+            undefined,
+            undefined,
+            undefined,
+            factory.createIdentifier('event'),
+            undefined,
+            factory.createTypeReferenceNode(factory.createIdentifier('SyntheticEvent'), undefined),
+            undefined,
+          ),
+        ],
+        undefined,
+        factory.createToken(ts.SyntaxKind.EqualsGreaterThanToken),
+        factory.createBlock(
+          [
+            factory.createExpressionStatement(
+              factory.createCallExpression(factory.createIdentifier(setStateName), undefined, [
+                factory.createPropertyAccessExpression(
+                  factory.createPropertyAccessExpression(
+                    factory.createIdentifier('event'),
+                    factory.createIdentifier('target'),
+                  ),
+                  factory.createIdentifier(targetFieldName),
+                ),
+              ]),
+            ),
+          ],
+          false,
+        ),
+      ),
+    );
+  };
+}
+
+function numericValueCallback(setStateName: string): JsxExpression {
   return factory.createJsxExpression(
     undefined,
     factory.createArrowFunction(
@@ -118,54 +161,33 @@ function syntheticEventTargetValueCallbackGenerator(stateName: string): JsxExpre
           undefined,
           undefined,
           undefined,
-          factory.createIdentifier('event'),
+          factory.createIdentifier('value'),
           undefined,
-          factory.createTypeReferenceNode(factory.createIdentifier('SyntheticEvent'), undefined),
+          factory.createKeywordTypeNode(ts.SyntaxKind.NumberKeyword),
           undefined,
         ),
       ],
       undefined,
       factory.createToken(ts.SyntaxKind.EqualsGreaterThanToken),
-      factory.createBlock(
-        [
-          factory.createExpressionStatement(
-            factory.createCallExpression(factory.createIdentifier(stateName), undefined, [
-              factory.createPropertyAccessExpression(
-                factory.createPropertyAccessExpression(
-                  factory.createIdentifier('event'),
-                  factory.createIdentifier('target'),
-                ),
-                factory.createIdentifier('value'),
-              ),
-            ]),
-          ),
-        ],
-        false,
-      ),
+      factory.createCallExpression(factory.createIdentifier(setStateName), undefined, [
+        factory.createIdentifier('value'),
+      ]),
     ),
   );
 }
 
-function numericValueCallbackGenerator(stateName: string): JsxExpression {
+function toggleBooleanStateCallback(setStateName: string, stateName: string): JsxExpression {
   return factory.createJsxExpression(
     undefined,
     factory.createArrowFunction(
       undefined,
       undefined,
-      [
-        factory.createParameterDeclaration(
-          undefined,
-          undefined,
-          undefined,
-          factory.createIdentifier('value'),
-          undefined,
-          factory.createKeywordTypeNode(ts.SyntaxKind.NumberKeyword),
-          undefined,
-        ),
-      ],
+      [],
       undefined,
       factory.createToken(ts.SyntaxKind.EqualsGreaterThanToken),
-      factory.createCallExpression(factory.createIdentifier(stateName), undefined, [factory.createIdentifier('value')]),
+      factory.createCallExpression(factory.createIdentifier(setStateName), undefined, [
+        factory.createPrefixUnaryExpression(ts.SyntaxKind.ExclamationToken, factory.createIdentifier(stateName)),
+      ]),
     ),
   );
 }
@@ -173,6 +195,7 @@ function numericValueCallbackGenerator(stateName: string): JsxExpression {
 // TODO: Update in here so we can support customer `change` events for form elements.
 export function buildOpeningElementControlEvents(
   componentType: string,
+  setStateName: string,
   stateName: string,
   event: string,
 ): JsxAttribute {
@@ -180,10 +203,10 @@ export function buildOpeningElementControlEvents(
   const controlEventBuilder =
     implementationOverrides && implementationOverrides[event]
       ? implementationOverrides[event]
-      : syntheticEventTargetValueCallbackGenerator;
+      : buildSyntheticEventTargetCallback('value');
   return factory.createJsxAttribute(
     factory.createIdentifier(mapGenericEventToReact(componentType as Primitive, event as StudioGenericEvent)),
-    controlEventBuilder(stateName),
+    controlEventBuilder(setStateName, stateName),
   );
 }
 
diff --git a/packages/codegen-ui/example-schemas/workflow/checkboxControlledElement.json b/packages/codegen-ui/example-schemas/workflow/checkboxControlledElement.json
new file mode 100644
index 000000000..a9a16d121
--- /dev/null
+++ b/packages/codegen-ui/example-schemas/workflow/checkboxControlledElement.json
@@ -0,0 +1,24 @@
+{
+    "id": "1234-5678-9010",
+    "componentType": "Flex",
+    "name": "CheckboxControlledElement",
+    "properties": {},
+    "children": [
+      {
+        "componentType": "CheckboxField",
+        "name": "Input",
+        "properties": {}
+      },
+      {
+        "componentType": "Text",
+        "name": "CheckboxFieldValue",
+        "properties": {
+          "label": {
+            "componentName": "Input",
+            "property": "checked"
+          }
+        }
+      }
+    ]
+  }
+  
\ No newline at end of file
diff --git a/packages/codegen-ui/example-schemas/workflow/switchControlledElement.json b/packages/codegen-ui/example-schemas/workflow/switchControlledElement.json
new file mode 100644
index 000000000..34738ef8c
--- /dev/null
+++ b/packages/codegen-ui/example-schemas/workflow/switchControlledElement.json
@@ -0,0 +1,24 @@
+{
+    "id": "1234-5678-9010",
+    "componentType": "Flex",
+    "name": "SwitchControlledElement",
+    "properties": {},
+    "children": [
+      {
+        "componentType": "SwitchField",
+        "name": "Input",
+        "properties": {}
+      },
+      {
+        "componentType": "Text",
+        "name": "SwitchFieldValue",
+        "properties": {
+          "label": {
+            "componentName": "Input",
+            "property": "isChecked"
+          }
+        }
+      }
+    ]
+  }
+  
\ No newline at end of file
diff --git a/packages/codegen-ui/example-schemas/workflow/twoWayBindings.json b/packages/codegen-ui/example-schemas/workflow/twoWayBindings.json
index ec300634b..014d17add 100644
--- a/packages/codegen-ui/example-schemas/workflow/twoWayBindings.json
+++ b/packages/codegen-ui/example-schemas/workflow/twoWayBindings.json
@@ -715,7 +715,7 @@
           "properties": {
             "label": {
               "componentName": "SwitchFieldInput",
-              "property": "value"
+              "property": "isChecked"
             }
           }
         },
@@ -733,7 +733,7 @@
               "parameters": {
                 "state": {
                   "componentName": "SwitchFieldInput",
-                  "property": "value",
+                  "property": "isChecked",
                   "set": {
                     "value": true
                   }
diff --git a/packages/test-generator/lib/components/workflow/twoWayBindings.json b/packages/test-generator/lib/components/workflow/twoWayBindings.json
index ec300634b..014d17add 100644
--- a/packages/test-generator/lib/components/workflow/twoWayBindings.json
+++ b/packages/test-generator/lib/components/workflow/twoWayBindings.json
@@ -715,7 +715,7 @@
           "properties": {
             "label": {
               "componentName": "SwitchFieldInput",
-              "property": "value"
+              "property": "isChecked"
             }
           }
         },
@@ -733,7 +733,7 @@
               "parameters": {
                 "state": {
                   "componentName": "SwitchFieldInput",
-                  "property": "value",
+                  "property": "isChecked",
                   "set": {
                     "value": true
                   }