fix: get grand child from component tree (#394)

* fix: get grand child from component tree

* test: add snapshot

packages/codegen-ui-react/lib/__tests__/workflow/__snapshots__/mutation.test.ts.snap

@@ -12,6 +12,8 @@ Array [
 ]
 `;
 
+exports[`getComponentFromComponentTree not found 1`] = `"Component NotFoundComponent not found in component tree Component"`;
+
 exports[`getComponentStateReferences basic 1`] = `
 Array [
   Object {

packages/codegen-ui-react/lib/__tests__/workflow/mutation.test.ts

@@ -14,7 +14,11 @@
   limitations under the License.
  */
 import { MutationAction, DataStoreUpdateItemAction } from '@aws-amplify/codegen-ui';
-import { getComponentStateReferences, getActionStateParameters } from '../../workflow/mutation';
+import {
+  getComponentStateReferences,
+  getActionStateParameters,
+  getComponentFromComponentTree,
+} from '../../workflow/mutation';
 
 describe('getComponentStateReferences', () => {
   test('basic', () => {
@@ -93,3 +97,39 @@ describe('getActionStateParameters', () => {
     expect(getActionStateParameters(action)).toMatchSnapshot();
   });
 });
+
+describe('getComponentFromComponentTree', () => {
+  const grandChildComponent = {
+    componentType: 'TextField',
+    name: 'GrandChildComponent',
+    properties: {},
+  };
+  const childComponent = {
+    componentType: 'Flex',
+    name: 'ChildComponent',
+    properties: {},
+    children: [grandChildComponent],
+  };
+  const component = {
+    componentType: 'Flex',
+    name: 'Component',
+    properties: {},
+    bindingProperties: {},
+    children: [childComponent],
+  };
+  test('same as root component', () => {
+    expect(getComponentFromComponentTree(component, 'Component')).toEqual(component);
+  });
+
+  test('child component', () => {
+    expect(getComponentFromComponentTree(component, 'ChildComponent')).toEqual(childComponent);
+  });
+
+  test('grandchild component', () => {
+    expect(getComponentFromComponentTree(component, 'GrandChildComponent')).toEqual(grandChildComponent);
+  });
+
+  test('not found', () => {
+    expect(() => getComponentFromComponentTree(component, 'NotFoundComponent')).toThrowErrorMatchingSnapshot();
+  });
+});

packages/codegen-ui-react/lib/workflow/mutation.ts

@@ -265,9 +265,7 @@ export function getComponentFromComponentTree(
     }
 
     if (currentComponent.children) {
-      return currentComponent.children.find(
-        (child: StudioComponentChild) => getComponentFromComponentTreeHelper(child) !== undefined,
-      );
+      return currentComponent.children.map(getComponentFromComponentTreeHelper).find((child) => child !== undefined);
     }
 
     return undefined;