Ensure that account tracker does not set balances incorrectly when us… (#17021)

* Ensure that account tracker does not set balances incorrectly when useMultiAccountBalanceChecker is false

* Lint fix

* Fix account setting in _updateAccount in useMultiAccountBalanceChecker === true case

* Fix _updateAccount

* Add unit tests

app/scripts/lib/account-tracker.js

@@ -93,6 +93,22 @@ export default class AccountTracker {
         }
       }, this.onboardingController.store.getState()),
     );
+
+    this.preferencesController.store.subscribe(
+      previousValueComparator(async (prevState, currState) => {
+        const { selectedAddress: prevSelectedAddress } = prevState;
+        const {
+          selectedAddress: currSelectedAddress,
+          useMultiAccountBalanceChecker,
+        } = currState;
+        if (
+          prevSelectedAddress !== currSelectedAddress &&
+          !useMultiAccountBalanceChecker
+        ) {
+          this._updateAccounts();
+        }
+      }, this.onboardingController.store.getState()),
+    );
   }
 
   start() {
@@ -321,6 +337,9 @@ export default class AccountTracker {
    * @returns {Promise} after the account balance is updated
    */
   async _updateAccount(address) {
+    const { useMultiAccountBalanceChecker } =
+      this.preferencesController.store.getState();
+
     let balance = '0x0';
 
     // query balance
@@ -339,8 +358,23 @@ export default class AccountTracker {
     if (!accounts[address]) {
       return;
     }
-    accounts[address] = result;
-    this.store.updateState({ accounts });
+
+    let newAccounts = accounts;
+    if (!useMultiAccountBalanceChecker) {
+      newAccounts = {};
+      Object.keys(accounts).forEach((accountAddress) => {
+        if (address !== accountAddress) {
+          newAccounts[accountAddress] = {
+            address: accountAddress,
+            balance: null,
+          };
+        }
+      });
+    }
+
+    newAccounts[address] = result;
+
+    this.store.updateState({ accounts: newAccounts });
   }
 
   /**
@@ -351,6 +385,12 @@ export default class AccountTracker {
    */
   async _updateAccountsViaBalanceChecker(addresses, deployedContractAddress) {
     const { accounts } = this.store.getState();
+    const newAccounts = {};
+    Object.keys(accounts).forEach((address) => {
+      if (!addresses.includes(address)) {
+        newAccounts[address] = { address, balance: null };
+      }
+    });
     this.ethersProvider = new ethers.providers.Web3Provider(this._provider);
 
     const ethContract = await new ethers.Contract(
@@ -365,9 +405,9 @@ export default class AccountTracker {
 
       addresses.forEach((address, index) => {
         const balance = balances[index] ? balances[index].toHexString() : '0x0';
-        accounts[address] = { address, balance };
+        newAccounts[address] = { address, balance };
       });
-      this.store.updateState({ accounts });
+      this.store.updateState({ accounts: newAccounts });
     } catch (error) {
       log.warn(
         `MetaMask - Account Tracker single call balance fetch failed`,

app/scripts/lib/account-tracker.test.js

@@ -0,0 +1,198 @@
+import { strict as assert } from 'assert';
+import EventEmitter from 'events';
+
+import { SINGLE_CALL_BALANCES_ADDRESS } from '../constants/contracts';
+
+import { createTestProviderTools } from '../../../test/stub/provider';
+import AccountTracker from './account-tracker';
+
+const noop = () => true;
+const currentNetworkId = '5';
+const VALID_ADDRESS = '0x0000000000000000000000000000000000000000';
+const VALID_ADDRESS_TWO = '0x0000000000000000000000000000000000000001';
+
+const INITIAL_BALANCE_1 = '0x1';
+const INITIAL_BALANCE_2 = '0x2';
+const UPDATE_BALANCE = '0xabc';
+
+// The below three values were generated by running MetaMask in the browser
+// The response to eth_call, which is called via `ethContract.balances`
+// in `_updateAccountsViaBalanceChecker` of account-tracker.js, needs to be properly
+// formatted or else ethers will throw an error.
+const ETHERS_CONTRACT_BALANCES_ETH_CALL_RETURN =
+  '0x0000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000038d7ea4c6800600000000000000000000000000000000000000000000000000000000000186a0';
+const EXPECTED_CONTRACT_BALANCE_1 = '0x038d7ea4c68006';
+const EXPECTED_CONTRACT_BALANCE_2 = '0x0186a0';
+
+const mockAccounts = {
+  [VALID_ADDRESS]: { address: VALID_ADDRESS, balance: INITIAL_BALANCE_1 },
+  [VALID_ADDRESS_TWO]: {
+    address: VALID_ADDRESS_TWO,
+    balance: INITIAL_BALANCE_2,
+  },
+};
+
+describe('Account Tracker', () => {
+  let provider,
+    blockTrackerStub,
+    providerResultStub,
+    useMultiAccountBalanceChecker,
+    accountTracker;
+
+  beforeEach(() => {
+    providerResultStub = {
+      eth_getBalance: UPDATE_BALANCE,
+      eth_call: ETHERS_CONTRACT_BALANCES_ETH_CALL_RETURN,
+    };
+    provider = createTestProviderTools({
+      scaffold: providerResultStub,
+      networkId: currentNetworkId,
+      chainId: currentNetworkId,
+    }).provider;
+
+    blockTrackerStub = new EventEmitter();
+    blockTrackerStub.getCurrentBlock = noop;
+    blockTrackerStub.getLatestBlock = noop;
+
+    accountTracker = new AccountTracker({
+      provider,
+      blockTracker: blockTrackerStub,
+      preferencesController: {
+        store: {
+          getState: () => ({
+            useMultiAccountBalanceChecker,
+          }),
+          subscribe: noop,
+        },
+      },
+      onboardingController: {
+        store: {
+          subscribe: noop,
+          getState: noop,
+        },
+      },
+    });
+  });
+
+  describe('_updateAccount', () => {
+    it('should update the passed address account balance, and leave other balances unchanged, if useMultiAccountBalanceChecker is true', async () => {
+      useMultiAccountBalanceChecker = true;
+      accountTracker.store.updateState({ accounts: { ...mockAccounts } });
+
+      await accountTracker._updateAccount(VALID_ADDRESS);
+
+      const newAccounts = accountTracker.store.getState();
+
+      const expectedAccounts = {
+        accounts: {
+          [VALID_ADDRESS]: { address: VALID_ADDRESS, balance: UPDATE_BALANCE },
+          [VALID_ADDRESS_TWO]: {
+            address: VALID_ADDRESS_TWO,
+            balance: INITIAL_BALANCE_2,
+          },
+        },
+        currentBlockGasLimit: '',
+      };
+
+      assert.deepEqual(newAccounts, expectedAccounts);
+    });
+
+    it('should not change accounts if the passed address is not in accounts', async () => {
+      accountTracker.store.updateState({ accounts: { ...mockAccounts } });
+
+      await accountTracker._updateAccount('fake address');
+
+      const newAccounts = accountTracker.store.getState();
+
+      const expectedAccounts = {
+        accounts: {
+          [VALID_ADDRESS]: {
+            address: VALID_ADDRESS,
+            balance: INITIAL_BALANCE_1,
+          },
+          [VALID_ADDRESS_TWO]: {
+            address: VALID_ADDRESS_TWO,
+            balance: INITIAL_BALANCE_2,
+          },
+        },
+        currentBlockGasLimit: '',
+      };
+
+      assert.deepEqual(newAccounts, expectedAccounts);
+    });
+
+    it('should update the passed address account balance, and set other balances to null, if useMultiAccountBalanceChecker is false', async () => {
+      useMultiAccountBalanceChecker = false;
+      accountTracker.store.updateState({ accounts: { ...mockAccounts } });
+
+      await accountTracker._updateAccount(VALID_ADDRESS);
+
+      const newAccounts = accountTracker.store.getState();
+
+      const expectedAccounts = {
+        accounts: {
+          [VALID_ADDRESS]: { address: VALID_ADDRESS, balance: UPDATE_BALANCE },
+          [VALID_ADDRESS_TWO]: { address: VALID_ADDRESS_TWO, balance: null },
+        },
+        currentBlockGasLimit: '',
+      };
+
+      assert.deepEqual(newAccounts, expectedAccounts);
+    });
+  });
+
+  describe('_updateAccountsViaBalanceChecker', () => {
+    it('should update the passed address account balance, and set other balances to null, if useMultiAccountBalanceChecker is false', async () => {
+      useMultiAccountBalanceChecker = true;
+      accountTracker.store.updateState({ accounts: { ...mockAccounts } });
+
+      await accountTracker._updateAccountsViaBalanceChecker(
+        [VALID_ADDRESS],
+        SINGLE_CALL_BALANCES_ADDRESS,
+      );
+
+      const newAccounts = accountTracker.store.getState();
+
+      const expectedAccounts = {
+        accounts: {
+          [VALID_ADDRESS]: {
+            address: VALID_ADDRESS,
+            balance: EXPECTED_CONTRACT_BALANCE_1,
+          },
+          [VALID_ADDRESS_TWO]: { address: VALID_ADDRESS_TWO, balance: null },
+        },
+        currentBlockGasLimit: '',
+      };
+
+      assert.deepEqual(newAccounts, expectedAccounts);
+    });
+
+    it('should update all balances if useMultiAccountBalanceChecker is true', async () => {
+      useMultiAccountBalanceChecker = true;
+      accountTracker.store.updateState({ accounts: { ...mockAccounts } });
+
+      await accountTracker._updateAccountsViaBalanceChecker(
+        [VALID_ADDRESS, VALID_ADDRESS_TWO],
+        SINGLE_CALL_BALANCES_ADDRESS,
+      );
+
+      const newAccounts = accountTracker.store.getState();
+
+      const expectedAccounts = {
+        accounts: {
+          [VALID_ADDRESS]: {
+            address: VALID_ADDRESS,
+            balance: EXPECTED_CONTRACT_BALANCE_1,
+          },
+          [VALID_ADDRESS_TWO]: {
+            address: VALID_ADDRESS_TWO,
+            balance: EXPECTED_CONTRACT_BALANCE_2,
+          },
+        },
+        currentBlockGasLimit: '',
+      };
+
+      assert.deepEqual(newAccounts, expectedAccounts);
+    });
+  });
+});