From adf26f43fe6823ba15480a81a2513bb2a98bd7b2 Mon Sep 17 00:00:00 2001
From: Justin Shih <36183898+Jshhhh@users.noreply.github.com>
Date: Tue, 23 May 2023 11:08:14 -0700
Subject: [PATCH] fix: fix storage field prop import to storagemanager (#1016)

Co-authored-by: Justin Shih <jushih@amazon.com>
---
 ...studio-ui-codegen-react-forms.test.ts.snap | 46 +++++++++++++++++--
 .../studio-ui-codegen-react-forms.test.ts     |  3 +-
 .../forms/form-renderer-helper/type-helper.ts |  9 +++-
 3 files changed, 52 insertions(+), 6 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 8ea161854..c9eecf0b4 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
@@ -24848,10 +24848,10 @@ exports[`amplify form renderer tests forms with StorageField tests should render
 "/* eslint-disable */
 import * as React from \\"react\\";
 import { Button, Flex, Grid, TextField } from \\"@aws-amplify/ui-react\\";
+import { StorageManager } from \\"@aws-amplify/ui-react-storage\\";
 import { Field, getOverrideProps } from \\"@aws-amplify/ui-react/internal\\";
 import { Product } from \\"../models\\";
 import { fetchByPath, processFile, validateField } from \\"./utils\\";
-import { StorageManager } from \\"@aws-amplify/ui-react-storage\\";
 import { DataStore } from \\"aws-amplify\\";
 export default function CreateProductForm(props) {
   const {
@@ -25059,10 +25059,10 @@ exports[`amplify form renderer tests forms with StorageField tests should render
 "/* eslint-disable */
 import * as React from \\"react\\";
 import { Button, Flex, Grid, TextField } from \\"@aws-amplify/ui-react\\";
+import { StorageManager } from \\"@aws-amplify/ui-react-storage\\";
 import { Field, getOverrideProps } from \\"@aws-amplify/ui-react/internal\\";
 import { Product } from \\"../models\\";
 import { fetchByPath, processFile, validateField } from \\"./utils\\";
-import { StorageManager } from \\"@aws-amplify/ui-react-storage\\";
 import { DataStore } from \\"aws-amplify\\";
 export default function UpdateProductForm(props) {
   const {
@@ -25287,14 +25287,54 @@ export default function UpdateProductForm(props) {
 "
 `;
 
+exports[`amplify form renderer tests forms with StorageField tests should render a update form with StorageField 2`] = `
+"import * as React from \\"react\\";
+import { GridProps, TextFieldProps } from \\"@aws-amplify/ui-react\\";
+import { StorageManagerProps } from \\"@aws-amplify/ui-react-storage\\";
+import { EscapeHatchProps } from \\"@aws-amplify/ui-react/internal\\";
+import { Product } from \\"../models\\";
+export declare type ValidationResponse = {
+    hasError: boolean;
+    errorMessage?: string;
+};
+export declare type ValidationFunction<T> = (value: T, validationResponse: ValidationResponse) => ValidationResponse | Promise<ValidationResponse>;
+export declare type UpdateProductFormInputValues = {
+    name?: string;
+    imgKeys?: string[];
+};
+export declare type UpdateProductFormValidationValues = {
+    name?: ValidationFunction<string>;
+    imgKeys?: ValidationFunction<string>;
+};
+export declare type PrimitiveOverrideProps<T> = Partial<T> & React.DOMAttributes<HTMLDivElement>;
+export declare type UpdateProductFormOverridesProps = {
+    UpdateProductFormGrid?: PrimitiveOverrideProps<GridProps>;
+    name?: PrimitiveOverrideProps<TextFieldProps>;
+    imgKeys?: PrimitiveOverrideProps<StorageManagerProps>;
+} & EscapeHatchProps;
+export declare type UpdateProductFormProps = React.PropsWithChildren<{
+    overrides?: UpdateProductFormOverridesProps | undefined | null;
+} & {
+    id?: string;
+    product?: Product;
+    onSubmit?: (fields: UpdateProductFormInputValues) => UpdateProductFormInputValues;
+    onSuccess?: (fields: UpdateProductFormInputValues) => void;
+    onError?: (fields: UpdateProductFormInputValues, errorMessage: string) => void;
+    onChange?: (fields: UpdateProductFormInputValues) => UpdateProductFormInputValues;
+    onValidate?: UpdateProductFormValidationValues;
+} & React.CSSProperties>;
+export default function UpdateProductForm(props: UpdateProductFormProps): React.ReactElement;
+"
+`;
+
 exports[`amplify form renderer tests forms with StorageField tests should render a update form with StorageField on non-array field 1`] = `
 "/* eslint-disable */
 import * as React from \\"react\\";
 import { Button, Flex, Grid, TextField } from \\"@aws-amplify/ui-react\\";
+import { StorageManager } from \\"@aws-amplify/ui-react-storage\\";
 import { Field, getOverrideProps } from \\"@aws-amplify/ui-react/internal\\";
 import { Product } from \\"../models\\";
 import { fetchByPath, processFile, validateField } from \\"./utils\\";
-import { StorageManager } from \\"@aws-amplify/ui-react-storage\\";
 import { DataStore } from \\"aws-amplify\\";
 export default function UpdateProductForm(props) {
   const {
diff --git a/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react-forms.test.ts b/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react-forms.test.ts
index 665ddfd7c..82fe46d5a 100644
--- a/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react-forms.test.ts
+++ b/packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react-forms.test.ts
@@ -663,12 +663,13 @@ describe('amplify form renderer tests', () => {
     });
 
     it('should render a update form with StorageField', () => {
-      const { componentText } = generateWithAmplifyFormRenderer(
+      const { componentText, declaration } = generateWithAmplifyFormRenderer(
         'forms/product-datastore-update',
         'datastore/product',
         undefined,
       );
       expect(componentText).toMatchSnapshot();
+      expect(declaration).toMatchSnapshot();
     });
   });
 
diff --git a/packages/codegen-ui-react/lib/forms/form-renderer-helper/type-helper.ts b/packages/codegen-ui-react/lib/forms/form-renderer-helper/type-helper.ts
index 958931bc3..33cd50976 100644
--- a/packages/codegen-ui-react/lib/forms/form-renderer-helper/type-helper.ts
+++ b/packages/codegen-ui-react/lib/forms/form-renderer-helper/type-helper.ts
@@ -493,7 +493,13 @@ export const buildOverrideTypesBindings = (
       if (field.split('.').length > 1 || !isValidVariableName(field)) {
         propKey = factory.createStringLiteral(field);
       }
-      const componentTypePropName = `${formDefinition.elements[field].componentType}Props`;
+      let componentTypePropName = `${formDefinition.elements[field].componentType}Props`;
+      if (formDefinition.elements[field].componentType === 'StorageField') {
+        componentTypePropName = 'StorageManagerProps';
+        importCollection.addImport(ImportSource.REACT_STORAGE, componentTypePropName);
+      } else {
+        importCollection.addImport(ImportSource.UI_REACT, componentTypePropName);
+      }
       typeNodes.push(
         factory.createPropertySignature(
           undefined,
@@ -504,7 +510,6 @@ export const buildOverrideTypesBindings = (
           ]),
         ),
       );
-      importCollection.addImport(ImportSource.UI_REACT, componentTypePropName);
     });
   });