From a1809ba558c9f10158e73ef0ac8c4b56f547a7ef Mon Sep 17 00:00:00 2001
From: Dane Pilcher <dppilche@amazon.com>
Date: Thu, 10 Feb 2022 17:13:00 +0000
Subject: [PATCH] fix: add state for state reference in set state param

---
 .../studio-ui-codegen-react.test.ts.snap      | 54 +++++++++++++++++++
 .../__tests__/studio-ui-codegen-react.test.ts |  4 ++
 .../codegen-ui-react/lib/workflow/mutation.ts |  5 +-
 .../workflow/nestedMutation.json              | 50 +++++++++++++++++
 4 files changed, 112 insertions(+), 1 deletion(-)
 create mode 100644 packages/codegen-ui/example-schemas/workflow/nestedMutation.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 68ae1eb87..aa50a182f 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
@@ -5776,6 +5776,60 @@ export default function SetStateWithoutInitialValue(
 }
 `;
 
+exports[`amplify render tests mutations supports nested mutation 1`] = `
+Object {
+  "componentText": "/* eslint-disable */
+import React from \\"react\\";
+import {
+  EscapeHatchProps,
+  getOverrideProps,
+  useStateMutationAction,
+} from \\"@aws-amplify/ui-react/internal\\";
+import { Button, Flex, FlexProps, Text } from \\"@aws-amplify/ui-react\\";
+
+export type NestedMutationProps = React.PropsWithChildren<
+  Partial<FlexProps> & {
+    overrides?: EscapeHatchProps | undefined | null;
+  }
+>;
+export default function NestedMutation(
+  props: NestedMutationProps
+): React.ReactElement {
+  const { overrides, ...rest } = props;
+  const [mutatedValueChildren, setMutatedValueChildren] =
+    useStateMutationAction(\\"Default Value\\");
+  const [stateSourceChildren, setStateSourceChildren] =
+    useStateMutationAction(\\"State Value\\");
+  const statePropertyMutationClick = () => {
+    setMutatedValueChildren(stateSourceLabel);
+  };
+  return (
+    /* @ts-ignore: TS2322 */
+    <Flex {...rest} {...getOverrideProps(overrides, \\"NestedMutation\\")}>
+      <Text
+        children={mutatedValueChildren}
+        {...getOverrideProps(overrides, \\"MutatedValue\\")}
+      ></Text>
+      <Button
+        children=\\"Apply State Property Mutation\\"
+        onClick={() => {
+          statePropertyMutationClick();
+        }}
+        {...getOverrideProps(overrides, \\"StatePropertyMutation\\")}
+      ></Button>
+      <Text
+        children={stateSourceChildren}
+        {...getOverrideProps(overrides, \\"StateSource\\")}
+      ></Text>
+    </Flex>
+  );
+}
+",
+  "declaration": undefined,
+  "renderComponentToFilesystem": [Function],
+}
+`;
+
 exports[`amplify render tests mutations supports two-way data binding on form elements 1`] = `
 Object {
   "componentText": "/* eslint-disable */
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 e014c2e8b..81e5f9aae 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
@@ -377,6 +377,10 @@ describe('amplify render tests', () => {
     it('controls an input that is modified by a button', () => {
       expect(generateWithAmplifyRenderer('workflow/inputMutationOnClick')).toMatchSnapshot();
     });
+
+    it('supports nested mutation', () => {
+      expect(generateWithAmplifyRenderer('workflow/nestedMutation')).toMatchSnapshot();
+    });
   });
 
   describe('default value', () => {
diff --git a/packages/codegen-ui-react/lib/workflow/mutation.ts b/packages/codegen-ui-react/lib/workflow/mutation.ts
index fc2aab45f..148bed434 100644
--- a/packages/codegen-ui-react/lib/workflow/mutation.ts
+++ b/packages/codegen-ui-react/lib/workflow/mutation.ts
@@ -102,6 +102,9 @@ export function getActionStateParameters(action: ActionStudioComponentEvent): St
         if (key === 'fields' || key === 'attributes') {
           return Object.values(parameter);
         }
+        if (isSetStateParameter(parameter)) {
+          return [parameter, parameter.set];
+        }
         return parameter;
       })
       .filter((parameter) => isStateProperty(parameter) || isSetStateParameter(parameter));
@@ -326,7 +329,7 @@ function mutationReferenceReducerWithComponentType(componentType: string) {
     return {
       ...mutationReferences,
       [stateReference.property]: propertyReferences.concat([
-        { addControlEvent: PrimitivesWithChangeEvent.has(componentType as Primitive) || !('set' in stateReference) },
+        { addControlEvent: PrimitivesWithChangeEvent.has(componentType as Primitive) },
       ]),
     };
   };
diff --git a/packages/codegen-ui/example-schemas/workflow/nestedMutation.json b/packages/codegen-ui/example-schemas/workflow/nestedMutation.json
new file mode 100644
index 000000000..2dddaa1c4
--- /dev/null
+++ b/packages/codegen-ui/example-schemas/workflow/nestedMutation.json
@@ -0,0 +1,50 @@
+{
+  "id": "1234-5678-9010",
+  "componentType": "Flex",
+  "name": "NestedMutation",
+  "properties": {},
+  "children": [
+    {
+      "componentType": "Text",
+      "name": "MutatedValue",
+      "properties": {
+        "label": {
+          "value": "Default Value"
+        }
+      }
+    },
+    {
+      "componentType": "Button",
+      "name": "StatePropertyMutation",
+      "events": {
+        "click": {
+          "action": "Amplify.Mutation",
+          "parameters": {
+            "state": {
+              "componentName": "MutatedValue",
+              "property": "label",
+              "set": {
+                "componentName": "StateSource",
+                "property": "label"
+              }
+            }
+          }
+        }
+      },
+      "properties": {
+        "children": {
+          "value": "Apply State Property Mutation"
+        }
+      }
+    },
+    {
+      "componentType": "Text",
+      "name": "StateSource",
+      "properties": {
+        "label": {
+          "value": "State Value"
+        }
+      }
+    }
+  ]
+}