diff --git a/app/scripts/lib/account-tracker.js b/app/scripts/lib/account-tracker.js index ddbaf7ec4..6757dcf4a 100644 --- a/app/scripts/lib/account-tracker.js +++ b/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`, diff --git a/app/scripts/lib/account-tracker.test.js b/app/scripts/lib/account-tracker.test.js new file mode 100644 index 000000000..7320c71eb --- /dev/null +++ b/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); + }); + }); +});