Use fromAccount instead of selectedAcccount for account updates on edit (#15449)

Author: brad-decker

ui/ducks/send/send.js

@@ -1454,28 +1454,30 @@ const slice = createSlice({
   extraReducers: (builder) => {
     builder
       .addCase(ACCOUNT_CHANGED, (state, action) => {
-        // If we are on the edit flow then we need to watch for changes to the
-        // current account.address in state and keep balance updated
-        // appropriately
-        if (
-          state.stage === SEND_STAGES.EDIT &&
-          action.payload.account.address === state.selectedAccount.address
-        ) {
-          // This event occurs when the user's account details update due to
-          // background state changes. If the account that is being updated is
-          // the current from account on the edit flow we need to update
-          // the balance for the account and revalidate the send state.
-          state.selectedAccount.balance = action.payload.account.balance;
-          // We need to update the asset balance if the asset is the native
-          // network asset. Once we update the balance we recompute error state.
+        // This event occurs when the user's account details update due to
+        // background state changes. If the account that is being updated is
+        // the current from account on the edit flow we need to update
+        // the balance for the account and revalidate the send state.
+        if (state.stage === SEND_STAGES.EDIT && action.payload.account) {
           const draftTransaction =
             state.draftTransactions[state.currentTransactionUUID];
-          if (draftTransaction?.asset.type === ASSET_TYPES.NATIVE) {
-            draftTransaction.asset.balance = action.payload.account.balance;
+          if (
+            draftTransaction &&
+            draftTransaction.fromAccount &&
+            draftTransaction.fromAccount.address ===
+              action.payload.account.address
+          ) {
+            draftTransaction.fromAccount.balance =
+              action.payload.account.balance;
+            // We need to update the asset balance if the asset is the native
+            // network asset. Once we update the balance we recompute error state.
+            if (draftTransaction.asset.type === ASSET_TYPES.NATIVE) {
+              draftTransaction.asset.balance = action.payload.account.balance;
+            }
+            slice.caseReducers.validateAmountField(state);
+            slice.caseReducers.validateGasField(state);
+            slice.caseReducers.validateSendState(state);
           }
-          slice.caseReducers.validateAmountField(state);
-          slice.caseReducers.validateGasField(state);
-          slice.caseReducers.validateSendState(state);
         }
       })
       .addCase(ADDRESS_BOOK_UPDATED, (state, action) => {

ui/ducks/send/send.test.js

@@ -1154,9 +1154,14 @@ describe('Send Slice', () => {
     });
 
     describe('Account Changed', () => {
-      it('should', () => {
+      it('should correctly update the fromAccount in an edit', () => {
         const accountsChangedState = {
-          ...INITIAL_SEND_STATE_FOR_EXISTING_DRAFT,
+          ...getInitialSendStateWithExistingTxState({
+            fromAccount: {
+              address: '0xAddress',
+              balance: '0x0',
+            },
+          }),
           stage: SEND_STAGES.EDIT,
           selectedAccount: {
             address: '0xAddress',
@@ -1176,11 +1181,42 @@ describe('Send Slice', () => {
 
         const result = sendReducer(accountsChangedState, action);
 
-        expect(result.selectedAccount.balance).toStrictEqual(
+        const draft = getTestUUIDTx(result);
+
+        expect(draft.fromAccount.balance).toStrictEqual(
           action.payload.account.balance,
         );
       });
 
+      it('should gracefully handle missing account param in payload', () => {
+        const accountsChangedState = {
+          ...getInitialSendStateWithExistingTxState({
+            fromAccount: {
+              address: '0xAddress',
+              balance: '0x0',
+            },
+          }),
+          stage: SEND_STAGES.EDIT,
+          selectedAccount: {
+            address: '0xAddress',
+            balance: '0x0',
+          },
+        };
+
+        const action = {
+          type: 'ACCOUNT_CHANGED',
+          payload: {
+            account: undefined,
+          },
+        };
+
+        const result = sendReducer(accountsChangedState, action);
+
+        const draft = getTestUUIDTx(result);
+
+        expect(draft.fromAccount.balance).toStrictEqual('0x0');
+      });
+
       it(`should not edit account balance if action payload address is not the same as state's address`, () => {
         const accountsChangedState = {
           ...INITIAL_SEND_STATE_FOR_EXISTING_DRAFT,