From ffe765970e7e26579790c9e1a9fc006624c660dc Mon Sep 17 00:00:00 2001
From: Roshane Pascual <rtpascual@users.noreply.github.com>
Date: Thu, 8 Dec 2022 11:44:36 -0800
Subject: [PATCH] fix: update import file path to handle aliased children
 (#814)

---
 .../studio-ui-codegen-react.test.ts.snap      | 40 +++++++++++++++++++
 .../__tests__/studio-ui-codegen-react.test.ts | 15 +++++++
 .../amplify-ui-renderers/customComponent.ts   |  7 +++-
 .../example-schemas/custom/aliasParent.json   | 14 +++++++
 4 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 packages/codegen-ui/example-schemas/custom/aliasParent.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 944c8a7d0..c10043729 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
@@ -4878,6 +4878,46 @@ export default function CustomButton(
 "
 `;
 
+exports[`amplify render tests custom components alias parent should render declarations 1`] = `
+"import * as React from \\"react\\";
+import { EscapeHatchProps } from \\"@aws-amplify/ui-react/internal\\";
+import { ViewTestProps } from \\"./ViewTest\\";
+export declare type CustomParentAndChildrenProps = React.PropsWithChildren<Partial<ViewTestProps> & {
+    overrides?: EscapeHatchProps | undefined | null;
+}>;
+export default function CustomParentAndChildren(props: CustomParentAndChildrenProps): React.ReactElement;
+"
+`;
+
+exports[`amplify render tests custom components alias parent should render parent with aliased child instead of primitive 1`] = `
+"/* eslint-disable */
+import * as React from \\"react\\";
+import {
+  EscapeHatchProps,
+  getOverrideProps,
+} from \\"@aws-amplify/ui-react/internal\\";
+import { GridCustom } from \\"./Grid\\";
+import ViewTest, { ViewTestProps } from \\"./ViewTest\\";
+
+export type AliasParentProps = React.PropsWithChildren<
+  Partial<ViewTestProps> & {
+    overrides?: EscapeHatchProps | undefined | null;
+  }
+>;
+export default function AliasParent(
+  props: AliasParentProps
+): React.ReactElement {
+  const { overrides, ...rest } = props;
+  return (
+    /* @ts-ignore: TS2322 */
+    <ViewTest {...getOverrideProps(overrides, \\"AliasParent\\")} {...rest}>
+      <GridCustom {...getOverrideProps(overrides, \\"Grid\\")}></GridCustom>
+    </ViewTest>
+  );
+}
+"
+`;
+
 exports[`amplify render tests custom components custom children should render component with custom children 1`] = `
 "/* eslint-disable */
 import * as React from \\"react\\";
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 290864832..1dbeeb367 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
@@ -540,6 +540,21 @@ describe('amplify render tests', () => {
         ).toMatchSnapshot();
       });
     });
+
+    describe('alias parent', () => {
+      it('should render parent with aliased child instead of primitive', () => {
+        expect(generateWithAmplifyRenderer('custom/aliasParent').componentText).toMatchSnapshot();
+      });
+
+      it('should render declarations', () => {
+        expect(
+          generateWithAmplifyRenderer('custom/customParentAndChildren', {
+            script: ScriptKind.JS,
+            renderTypeDeclarations: true,
+          }).declaration,
+        ).toMatchSnapshot();
+      });
+    });
   });
 
   describe('primitives', () => {
diff --git a/packages/codegen-ui-react/lib/amplify-ui-renderers/customComponent.ts b/packages/codegen-ui-react/lib/amplify-ui-renderers/customComponent.ts
index ee7b1210d..77579febb 100644
--- a/packages/codegen-ui-react/lib/amplify-ui-renderers/customComponent.ts
+++ b/packages/codegen-ui-react/lib/amplify-ui-renderers/customComponent.ts
@@ -15,6 +15,7 @@
  */
 import { StudioComponentChild } from '@aws-amplify/codegen-ui';
 import { JsxChild, JsxElement, factory } from 'typescript';
+import { Primitive } from '../primitive';
 import { ReactComponentRenderer } from '../react-component-renderer';
 
 export default class CustomComponentRenderer<TPropIn> extends ReactComponentRenderer<TPropIn> {
@@ -26,7 +27,11 @@ export default class CustomComponentRenderer<TPropIn> extends ReactComponentRend
       factory.createJsxClosingElement(factory.createIdentifier(this.component.componentType)),
     );
 
-    this.importCollection.addImport(`./${this.component.componentType}`, 'default');
+    if (Object.keys(Primitive).includes(this.component.name) && this.component.componentType.match(/(Custom$)/g)) {
+      this.importCollection.addImport(`./${this.component.name}`, this.component.componentType);
+    } else {
+      this.importCollection.addImport(`./${this.component.componentType}`, 'default');
+    }
 
     return element;
   }
diff --git a/packages/codegen-ui/example-schemas/custom/aliasParent.json b/packages/codegen-ui/example-schemas/custom/aliasParent.json
new file mode 100644
index 000000000..dc75c398f
--- /dev/null
+++ b/packages/codegen-ui/example-schemas/custom/aliasParent.json
@@ -0,0 +1,14 @@
+{
+    "name": "AliasParent",
+    "children": [
+        {
+            "name": "Grid",
+            "componentType": "GridCustom",
+            "properties": {}
+        }
+    ],
+    "properties": {},
+    "id": "1234-1234",
+    "componentType": "ViewTest",
+    "schemaVersion": "1.0"
+}
\ No newline at end of file