From b98a8f452bfa658f1c58ebe6a5e75fece40d72d9 Mon Sep 17 00:00:00 2001
From: Alexander Harris <alharris@amazon.com>
Date: Mon, 31 Jan 2022 18:05:54 -0800
Subject: [PATCH] fix: add missing import statement, and correct controlled
 component state tracking

---
 .../studio-ui-codegen-react.test.ts.snap      |  6 +-
 .../lib/imports/import-mapping.ts             | 12 ++++
 .../lib/react-component-renderer.ts           | 11 ++--
 .../react-component-with-children-renderer.ts | 11 ++--
 .../lib/react-studio-template-renderer.ts     |  2 +-
 .../codegen-ui-react/lib/workflow/action.ts   | 24 ++++----
 .../codegen-ui-react/lib/workflow/mutation.ts | 56 ++++++++++++++++++-
 .../workflow/authSignOutAction.json           | 29 +++++-----
 .../workflow/dataStoreCreateItem.json         |  1 -
 .../workflow/dataStoreDeleteItem.json         |  1 -
 .../workflow/dataStoreUpdateItem.json         |  1 -
 11 files changed, 112 insertions(+), 42 deletions(-)

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 ac30ff607..c19e73fee 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
@@ -5167,6 +5167,7 @@ import {
   EscapeHatchProps,
   getOverrideProps,
   useDataStoreUpdateAction,
+  useStateMutationAction,
 } from \\"@aws-amplify/ui-react/internal\\";
 import { Button, Flex, FlexProps, TextField } from \\"@aws-amplify/ui-react\\";
 import { Customer } from \\"../models\\";
@@ -5191,7 +5192,9 @@ export default function MyForm(props: MyFormProps): React.ReactElement {
       <TextField
         label=\\"Username\\"
         value={usernameTextFieldValue}
-        onChange={setUsernameTextFieldValue}
+        onChange={(event: SyntheticEvent) => {
+          setUsernameTextFieldValue(event.target.value);
+        }}
         {...getOverrideProps(overrides, \\"UsernameTextField\\")}
       ></TextField>
       <Button
@@ -5217,6 +5220,7 @@ import React from \\"react\\";
 import {
   EscapeHatchProps,
   getOverrideProps,
+  useStateMutationAction,
 } from \\"@aws-amplify/ui-react/internal\\";
 import { Button, Flex, FlexProps } from \\"@aws-amplify/ui-react\\";
 
diff --git a/packages/codegen-ui-react/lib/imports/import-mapping.ts b/packages/codegen-ui-react/lib/imports/import-mapping.ts
index 474ff4a61..e5e37cd98 100644
--- a/packages/codegen-ui-react/lib/imports/import-mapping.ts
+++ b/packages/codegen-ui-react/lib/imports/import-mapping.ts
@@ -32,6 +32,12 @@ export enum ImportValue {
   USE_DATA_STORE_BINDING = 'useDataStoreBinding',
   CREATE_THEME = 'createTheme',
   MERGE_VARIANTS_OVERRIDES = 'mergeVariantsAndOverrides',
+  USE_NAVIGATE_ACTION = 'useNavigateAction',
+  USE_DATA_STORE_CREATE_ACTION = 'useDataStoreCreateAction',
+  USE_DATA_STORE_UPDATE_ACTION = 'useDataStoreUpdateAction',
+  USE_DATA_STORE_DELETE_ACTION = 'useDataStoreDeleteAction',
+  USE_AUTH_SIGN_OUT_ACTION = 'useAuthSignOutAction',
+  USE_STATE_MUTATION_ACTION = 'useStateMutationAction',
 }
 
 export const ImportMapping: Record<ImportValue, ImportSource> = {
@@ -46,4 +52,10 @@ export const ImportMapping: Record<ImportValue, ImportSource> = {
   [ImportValue.SORT_DIRECTION]: ImportSource.AMPLIFY_DATASTORE,
   [ImportValue.SORT_PREDICATE]: ImportSource.AMPLIFY_DATASTORE,
   [ImportValue.MERGE_VARIANTS_OVERRIDES]: ImportSource.UI_REACT_INTERNAL,
+  [ImportValue.USE_NAVIGATE_ACTION]: ImportSource.UI_REACT_INTERNAL,
+  [ImportValue.USE_DATA_STORE_CREATE_ACTION]: ImportSource.UI_REACT_INTERNAL,
+  [ImportValue.USE_DATA_STORE_UPDATE_ACTION]: ImportSource.UI_REACT_INTERNAL,
+  [ImportValue.USE_DATA_STORE_DELETE_ACTION]: ImportSource.UI_REACT_INTERNAL,
+  [ImportValue.USE_AUTH_SIGN_OUT_ACTION]: ImportSource.UI_REACT_INTERNAL,
+  [ImportValue.USE_STATE_MUTATION_ACTION]: ImportSource.UI_REACT_INTERNAL,
 };
diff --git a/packages/codegen-ui-react/lib/react-component-renderer.ts b/packages/codegen-ui-react/lib/react-component-renderer.ts
index 14f0b5080..93fee4642 100644
--- a/packages/codegen-ui-react/lib/react-component-renderer.ts
+++ b/packages/codegen-ui-react/lib/react-component-renderer.ts
@@ -34,7 +34,11 @@ import {
   getStateName,
   getSetStateName,
 } from './react-component-render-helper';
-import { buildOpeningElementEvents, filterStateReferencesForComponent } from './workflow';
+import {
+  buildOpeningElementControlEvents,
+  buildOpeningElementEvents,
+  filterStateReferencesForComponent,
+} from './workflow';
 import { ImportCollection, ImportSource, ImportValue } from './imports';
 
 export class ReactComponentRenderer<TPropIn> extends ComponentRendererBase<
@@ -78,10 +82,9 @@ export class ReactComponentRenderer<TPropIn> extends ComponentRendererBase<
     const controlEventAttributes = Object.entries(localStateReferences)
       .filter(([, { addControlEvent }]) => addControlEvent)
       .map(([key]) =>
-        buildOpeningElementEvents(
-          { bindingEvent: getSetStateName({ componentName: this.component.name || '', property: key }) },
+        buildOpeningElementControlEvents(
+          getSetStateName({ componentName: this.component.name || '', property: key }),
           'change', // TODO: use component event mapping
-          this.component.name,
         ),
       );
     const attributes = propertyAttributes.concat(eventAttributes).concat(controlEventAttributes);
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 bfa3d6d8f..666f0e3e7 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
@@ -28,7 +28,11 @@ import {
   getStateName,
   getSetStateName,
 } from './react-component-render-helper';
-import { buildOpeningElementEvents, filterStateReferencesForComponent } from './workflow';
+import {
+  buildOpeningElementControlEvents,
+  buildOpeningElementEvents,
+  filterStateReferencesForComponent,
+} from './workflow';
 import Primitive, { PrimitiveChildrenPropMapping } from './primitive';
 
 export class ReactComponentWithChildrenRenderer<TPropIn> extends ComponentWithChildrenRendererBase<
@@ -76,10 +80,9 @@ export class ReactComponentWithChildrenRenderer<TPropIn> extends ComponentWithCh
     const controlEventAttributes = Object.entries(localStateReferences)
       .filter(([, { addControlEvent }]) => addControlEvent)
       .map(([key]) =>
-        buildOpeningElementEvents(
-          { bindingEvent: getSetStateName({ componentName: this.component.name || '', property: key }) },
+        buildOpeningElementControlEvents(
+          getSetStateName({ componentName: this.component.name || '', property: key }),
           'change', // TODO: use component event mapping
-          this.component.name,
         ),
       );
     const attributes = propertyAttributes.concat(eventAttributes).concat(controlEventAttributes);
diff --git a/packages/codegen-ui-react/lib/react-studio-template-renderer.ts b/packages/codegen-ui-react/lib/react-studio-template-renderer.ts
index 7a65b554d..ffc109ef1 100644
--- a/packages/codegen-ui-react/lib/react-studio-template-renderer.ts
+++ b/packages/codegen-ui-react/lib/react-studio-template-renderer.ts
@@ -590,7 +590,7 @@ export abstract class ReactStudioTemplateRenderer extends StudioTemplateRenderer
       statements.push(this.buildOverridesFromVariantsAndProp());
     }
 
-    const stateStatements = buildStateStatements(component, this.stateReferences);
+    const stateStatements = buildStateStatements(component, this.stateReferences, this.importCollection);
     stateStatements.forEach((entry) => {
       statements.push(entry);
     });
diff --git a/packages/codegen-ui-react/lib/workflow/action.ts b/packages/codegen-ui-react/lib/workflow/action.ts
index ee66877b4..12ac2c1b1 100644
--- a/packages/codegen-ui-react/lib/workflow/action.ts
+++ b/packages/codegen-ui-react/lib/workflow/action.ts
@@ -24,7 +24,7 @@ import {
   MutationAction,
 } from '@aws-amplify/codegen-ui';
 import { isActionEvent, propertyToExpression, getSetStateName } from '../react-component-render-helper';
-import { ImportCollection, ImportSource } from '../imports';
+import { ImportCollection, ImportSource, ImportValue } from '../imports';
 
 enum Action {
   'Amplify.Navigate' = 'Amplify.Navigate',
@@ -37,13 +37,13 @@ enum Action {
 
 export default Action;
 
-export const ActionNameMapping: Partial<Record<Action, string>> = {
-  [Action['Amplify.Navigate']]: 'useNavigateAction',
-  [Action['Amplify.DataStoreCreateItem']]: 'useDataStoreCreateAction',
-  [Action['Amplify.DataStoreUpdateItem']]: 'useDataStoreUpdateAction',
-  [Action['Amplify.DataStoreDeleteItem']]: 'useDataStoreDeleteAction',
-  [Action['Amplify.AuthSignOut']]: 'useAuthSignOutAction',
-  [Action['Amplify.Mutation']]: 'useStateMutationAction',
+export const ActionNameMapping: Partial<Record<Action, ImportValue>> = {
+  [Action['Amplify.Navigate']]: ImportValue.USE_NAVIGATE_ACTION,
+  [Action['Amplify.DataStoreCreateItem']]: ImportValue.USE_DATA_STORE_CREATE_ACTION,
+  [Action['Amplify.DataStoreUpdateItem']]: ImportValue.USE_DATA_STORE_UPDATE_ACTION,
+  [Action['Amplify.DataStoreDeleteItem']]: ImportValue.USE_DATA_STORE_DELETE_ACTION,
+  [Action['Amplify.AuthSignOut']]: ImportValue.USE_AUTH_SIGN_OUT_ACTION,
+  [Action['Amplify.Mutation']]: ImportValue.USE_STATE_MUTATION_ACTION,
 };
 
 export function isAction(action: string): action is Action {
@@ -54,7 +54,7 @@ export function isMutationAction(action: ActionStudioComponentEvent): action is
   return (action.action as Action) === Action['Amplify.Mutation'];
 }
 
-export function getActionHookName(action: string): string {
+export function getActionHookImportValue(action: string): ImportValue {
   const actionName = ActionNameMapping[Action[action as Action]];
   if (actionName === undefined) {
     throw new InvalidInputError(`${action} is not a valid action.`);
@@ -95,8 +95,8 @@ export function buildUseActionStatement(
     return buildMutationActionStatement(action, identifier);
   }
 
-  const actionHookName = getActionHookName(action.action);
-  importCollection.addImport(ImportSource.UI_REACT_INTERNAL, actionHookName);
+  const actionHookImportValue = getActionHookImportValue(action.action);
+  importCollection.addMappedImport(actionHookImportValue);
   return factory.createVariableStatement(
     undefined,
     factory.createVariableDeclarationList(
@@ -105,7 +105,7 @@ export function buildUseActionStatement(
           factory.createIdentifier(identifier),
           undefined,
           undefined,
-          factory.createCallExpression(factory.createIdentifier(actionHookName), undefined, [
+          factory.createCallExpression(factory.createIdentifier(actionHookImportValue), undefined, [
             buildActionParameters(action, importCollection),
           ]),
         ),
diff --git a/packages/codegen-ui-react/lib/workflow/mutation.ts b/packages/codegen-ui-react/lib/workflow/mutation.ts
index aaae4c781..3f970571e 100644
--- a/packages/codegen-ui-react/lib/workflow/mutation.ts
+++ b/packages/codegen-ui-react/lib/workflow/mutation.ts
@@ -13,13 +13,14 @@
   See the License for the specific language governing permissions and
   limitations under the License.
  */
-import ts, { Statement, factory } from 'typescript';
+import ts, { Statement, factory, JsxAttribute } from 'typescript';
 import {
   StudioComponent,
   StudioComponentChild,
   StateStudioComponentProperty,
   ActionStudioComponentEvent,
   StateReference,
+  StudioGenericEvent,
 } from '@aws-amplify/codegen-ui';
 import {
   isActionEvent,
@@ -29,6 +30,8 @@ import {
   getStateName,
   getSetStateName,
 } from '../react-component-render-helper';
+import { ImportCollection, ImportValue } from '../imports';
+import { mapGenericEventToReact } from './events';
 
 export function getComponentStateReferences(component: StudioComponent | StudioComponentChild) {
   const stateReferences: StateReference[] = [];
@@ -75,7 +78,56 @@ export function getActionStateParameters(action: ActionStudioComponentEvent): St
   return [];
 }
 
-export function buildStateStatements(component: StudioComponent, stateReferences: StateReference[]): Statement[] {
+export function buildOpeningElementControlEvents(stateName: string, event: string): JsxAttribute {
+  return factory.createJsxAttribute(
+    factory.createIdentifier(mapGenericEventToReact(event as StudioGenericEvent)),
+    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(stateName), undefined, [
+                factory.createPropertyAccessExpression(
+                  factory.createPropertyAccessExpression(
+                    factory.createIdentifier('event'),
+                    factory.createIdentifier('target'),
+                  ),
+                  factory.createIdentifier('value'),
+                ),
+              ]),
+            ),
+          ],
+          false,
+        ),
+      ),
+    ),
+  );
+}
+
+export function buildStateStatements(
+  component: StudioComponent,
+  stateReferences: StateReference[],
+  importCollection: ImportCollection,
+): Statement[] {
+  if (stateReferences.length > 0) {
+    importCollection.addMappedImport(ImportValue.USE_STATE_MUTATION_ACTION);
+  }
   return stateReferences.map((stateReference) => {
     return factory.createVariableStatement(
       undefined,
diff --git a/packages/codegen-ui/example-schemas/workflow/authSignOutAction.json b/packages/codegen-ui/example-schemas/workflow/authSignOutAction.json
index ff7bebc90..c9ae86554 100644
--- a/packages/codegen-ui/example-schemas/workflow/authSignOutAction.json
+++ b/packages/codegen-ui/example-schemas/workflow/authSignOutAction.json
@@ -1,20 +1,19 @@
 {
-    "componentType": "Button",
-    "name": "SignOutButton",
-    "properties": {
-      "label": {
-        "value": "Sign out"
-      }
-    },
-    "events": {
-      "click": {
-        "action": "Amplify.AuthSignOut",
-        "parameters": {
-          "global": {
-            "value": false
-          }
+  "componentType": "Button",
+  "name": "SignOutButton",
+  "properties": {
+    "label": {
+      "value": "Sign out"
+    }
+  },
+  "events": {
+    "click": {
+      "action": "Amplify.AuthSignOut",
+      "parameters": {
+        "global": {
+          "value": false
         }
       }
     }
   }
-  
\ No newline at end of file
+}
diff --git a/packages/codegen-ui/example-schemas/workflow/dataStoreCreateItem.json b/packages/codegen-ui/example-schemas/workflow/dataStoreCreateItem.json
index 69feb5305..b6dddc949 100644
--- a/packages/codegen-ui/example-schemas/workflow/dataStoreCreateItem.json
+++ b/packages/codegen-ui/example-schemas/workflow/dataStoreCreateItem.json
@@ -1,7 +1,6 @@
 {
   "componentType": "Button",
   "name": "CreateCustomerButton",
-  "properties": {},
   "events": {
     "click": {
       "action": "Amplify.DataStoreCreateItem",
diff --git a/packages/codegen-ui/example-schemas/workflow/dataStoreDeleteItem.json b/packages/codegen-ui/example-schemas/workflow/dataStoreDeleteItem.json
index a32f849ac..4f1f0f22b 100644
--- a/packages/codegen-ui/example-schemas/workflow/dataStoreDeleteItem.json
+++ b/packages/codegen-ui/example-schemas/workflow/dataStoreDeleteItem.json
@@ -1,7 +1,6 @@
 {
   "componentType": "Button",
   "name": "DeleteCustomerButton",
-  "properties": {},
   "events": {
     "click": {
       "action": "Amplify.DataStoreDeleteItem",
diff --git a/packages/codegen-ui/example-schemas/workflow/dataStoreUpdateItem.json b/packages/codegen-ui/example-schemas/workflow/dataStoreUpdateItem.json
index 25470eede..bfb7b907c 100644
--- a/packages/codegen-ui/example-schemas/workflow/dataStoreUpdateItem.json
+++ b/packages/codegen-ui/example-schemas/workflow/dataStoreUpdateItem.json
@@ -1,7 +1,6 @@
 {
   "componentType": "Button",
   "name": "UpdateCustomerButton",
-  "properties": {},
   "events": {
     "click": {
       "action": "Amplify.DataStoreUpdateItem",