From ad0beda3f05b086bdd4347589b4f8bde243be07a Mon Sep 17 00:00:00 2001
From: Scott Young <scoyou@amazon.com>
Date: Fri, 18 Nov 2022 11:46:30 -0800
Subject: [PATCH] fix: update model used

---
 .../studio-ui-codegen-react-forms.test.ts.snap            | 2 +-
 .../codegen-ui-react/lib/amplify-ui-renderers/form.ts     | 2 +-
 .../codegen-ui-react/lib/forms/form-renderer-helper.ts    | 8 ++++++--
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
index 1910a4c36..397aad288 100644
--- a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
+++ b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
@@ -4471,7 +4471,7 @@ export default function MyFlexUpdateForm(props) {
         }
         try {
           await DataStore.save(
-            Flex0.copyOf(flex0Record, (updated) => {
+            Flex0.copyOf(flexRecord, (updated) => {
               Object.assign(updated, modelFields);
             })
           );
diff --git a/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts b/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts
index b1627cc41..c75c02a4d 100644
--- a/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts
+++ b/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts
@@ -127,7 +127,7 @@ export default class FormRenderer extends ReactComponentRenderer<BaseComponentPr
         factory.createTryStatement(
           factory.createBlock(
             [
-              ...buildDataStoreExpression(formActionType, importedModelName),
+              ...buildDataStoreExpression(formActionType, importedModelName, dataTypeName),
               // call onSuccess hook if it exists
               factory.createIfStatement(
                 factory.createIdentifier('onSuccess'),
diff --git a/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts b/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts
index e2a21d263..7ae321582 100644
--- a/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts
+++ b/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts
@@ -552,7 +552,11 @@ export const buildOnChangeStatement = (
   );
 };
 
-export const buildDataStoreExpression = (dataStoreActionType: 'update' | 'create', importedModelName: string) => {
+export const buildDataStoreExpression = (
+  dataStoreActionType: 'update' | 'create',
+  importedModelName: string,
+  dataTypeName: string,
+) => {
   if (dataStoreActionType === 'update') {
     return [
       factory.createExpressionStatement(
@@ -571,7 +575,7 @@ export const buildDataStoreExpression = (dataStoreActionType: 'update' | 'create
                 ),
                 undefined,
                 [
-                  factory.createIdentifier(`${lowerCaseFirst(importedModelName)}Record`),
+                  factory.createIdentifier(`${lowerCaseFirst(dataTypeName)}Record`),
                   factory.createArrowFunction(
                     undefined,
                     undefined,