cache balances by chain id (#10545)

app/scripts/controllers/cached-balances.js

@@ -3,7 +3,7 @@ import { ObservableStore } from '@metamask/obs-store';
 /**
  * @typedef {Object} CachedBalancesOptions
  * @property {Object} accountTracker An {@code AccountTracker} reference
- * @property {Function} getNetwork A function to get the current network
+ * @property {Function} getCurrentChainId A function to get the current chain id
  * @property {Object} initState The initial controller state
  */
 
@@ -18,10 +18,10 @@ export default class CachedBalancesController {
    * @param {CachedBalancesOptions} [opts] - Controller configuration parameters
    */
   constructor(opts = {}) {
-    const { accountTracker, getNetwork } = opts;
+    const { accountTracker, getCurrentChainId } = opts;
 
     this.accountTracker = accountTracker;
-    this.getNetwork = getNetwork;
+    this.getCurrentChainId = getCurrentChainId;
 
     const initState = { cachedBalances: {}, ...opts.initState };
     this.store = new ObservableStore(initState);
@@ -30,37 +30,37 @@ export default class CachedBalancesController {
   }
 
   /**
-   * Updates the cachedBalances property for the current network. Cached balances will be updated to those in the passed accounts
+   * Updates the cachedBalances property for the current chain. Cached balances will be updated to those in the passed accounts
    * if balances in the passed accounts are truthy.
    *
-   * @param {Object} obj - The the recently updated accounts object for the current network
+   * @param {Object} obj - The the recently updated accounts object for the current chain
    * @returns {Promise<void>}
    */
   async updateCachedBalances({ accounts }) {
-    const network = await this.getNetwork();
+    const chainId = this.getCurrentChainId();
     const balancesToCache = await this._generateBalancesToCache(
       accounts,
-      network,
+      chainId,
     );
     this.store.updateState({
       cachedBalances: balancesToCache,
     });
   }
 
-  _generateBalancesToCache(newAccounts, currentNetwork) {
+  _generateBalancesToCache(newAccounts, chainId) {
     const { cachedBalances } = this.store.getState();
-    const currentNetworkBalancesToCache = { ...cachedBalances[currentNetwork] };
+    const currentChainBalancesToCache = { ...cachedBalances[chainId] };
 
     Object.keys(newAccounts).forEach((accountID) => {
       const account = newAccounts[accountID];
 
       if (account.balance) {
-        currentNetworkBalancesToCache[accountID] = account.balance;
+        currentChainBalancesToCache[accountID] = account.balance;
       }
     });
     const balancesToCache = {
       ...cachedBalances,
-      [currentNetwork]: currentNetworkBalancesToCache,
+      [chainId]: currentChainBalancesToCache,
     };
 
     return balancesToCache;

app/scripts/metamask-controller.js

@@ -218,7 +218,7 @@ export default class MetamaskController extends EventEmitter {
 
     this.cachedBalancesController = new CachedBalancesController({
       accountTracker: this.accountTracker,
-      getNetwork: this.networkController.getNetworkState.bind(
+      getCurrentChainId: this.networkController.getCurrentChainId.bind(
         this.networkController,
       ),
       initState: initState.CachedBalancesController,

test/unit/app/controllers/cached-balances-test.js

@@ -1,12 +1,13 @@
 import assert from 'assert';
 import sinon from 'sinon';
 import CachedBalancesController from '../../../../app/scripts/controllers/cached-balances';
+import { KOVAN_CHAIN_ID } from '../../../../shared/constants/network';
 
 describe('CachedBalancesController', function () {
   describe('updateCachedBalances', function () {
     it('should update the cached balances', async function () {
       const controller = new CachedBalancesController({
-        getNetwork: () => Promise.resolve(17),
+        getCurrentChainId: () => KOVAN_CHAIN_ID,
         accountTracker: {
           store: {
             subscribe: () => undefined,
@@ -26,7 +27,7 @@ describe('CachedBalancesController', function () {
       assert.equal(controller._generateBalancesToCache.callCount, 1);
       assert.deepEqual(controller._generateBalancesToCache.args[0], [
         'mockAccounts',
-        17,
+        KOVAN_CHAIN_ID,
       ]);
       assert.equal(
         controller.store.getState().cachedBalances,
@@ -45,7 +46,7 @@ describe('CachedBalancesController', function () {
         },
         initState: {
           cachedBalances: {
-            17: {
+            [KOVAN_CHAIN_ID]: {
               a: '0x1',
               b: '0x2',
               c: '0x3',
@@ -65,11 +66,11 @@ describe('CachedBalancesController', function () {
           b: { balance: null },
           c: { balance: '0x5' },
         },
-        17,
+        KOVAN_CHAIN_ID,
       );
 
       assert.deepEqual(result, {
-        17: {
+        [KOVAN_CHAIN_ID]: {
           a: '0x4',
           b: '0x2',
           c: '0x5',
@@ -91,7 +92,7 @@ describe('CachedBalancesController', function () {
         },
         initState: {
           cachedBalances: {
-            17: {
+            [KOVAN_CHAIN_ID]: {
               a: '0x1',
               b: '0x2',
               c: '0x3',
@@ -110,7 +111,7 @@ describe('CachedBalancesController', function () {
       );
 
       assert.deepEqual(result, {
-        17: {
+        [KOVAN_CHAIN_ID]: {
           a: '0x1',
           b: '0x2',
           c: '0x3',
@@ -127,7 +128,7 @@ describe('CachedBalancesController', function () {
     it('should subscribe to the account tracker with the updateCachedBalances method', async function () {
       const subscribeSpy = sinon.spy();
       const controller = new CachedBalancesController({
-        getNetwork: () => Promise.resolve(17),
+        getCurrentChainId: () => KOVAN_CHAIN_ID,
         accountTracker: {
           store: {
             subscribe: subscribeSpy,

ui/app/components/app/alerts/unconnected-account-alert/tests/unconnected-account-alert.test.js

@@ -11,10 +11,9 @@ import { renderWithProvider } from '../../../../../../../test/lib/render-helpers
 
 import * as actions from '../../../../../store/actions';
 import UnconnectedAccountAlert from '..';
+import { KOVAN_CHAIN_ID } from '../../../../../../../shared/constants/network';
 
 describe('Unconnected Account Alert', function () {
-  const network = '123';
-
   const selectedAddress = '0xec1adf982415d2ef5ec55899b9bfb8bc0f29251b';
 
   const identities = {
@@ -40,7 +39,7 @@ describe('Unconnected Account Alert', function () {
   };
 
   const cachedBalances = {
-    123: {
+    [KOVAN_CHAIN_ID]: {
       '0x0dcd5d886577d5081b0c52e242ef29e70be3e7bc': '0x0',
       '0xec1adf982415d2ef5ec55899b9bfb8bc0f29251b': '0x0',
     },
@@ -58,12 +57,14 @@ describe('Unconnected Account Alert', function () {
 
   const mockState = {
     metamask: {
-      network,
       selectedAddress,
       identities,
       accounts,
       cachedBalances,
       keyrings,
+      provider: {
+        chainId: KOVAN_CHAIN_ID,
+      },
       permissionsHistory: {
         'https://test.dapp': {
           eth_accounts: {

ui/app/selectors/selectors.js

@@ -127,9 +127,16 @@ export function getMetaMaskAccountsRaw(state) {
 }
 
 export function getMetaMaskCachedBalances(state) {
+  const chainId = getCurrentChainId(state);
+
+  // Fallback to fetching cached balances from network id
+  // this can eventually be removed
   const network = getCurrentNetworkId(state);
 
-  return state.metamask.cachedBalances[network];
+  return (
+    state.metamask.cachedBalances[chainId] ??
+    state.metamask.cachedBalances[network]
+  );
 }
 
 /**
@@ -155,7 +162,7 @@ export function isBalanceCached(state) {
 }
 
 export function getSelectedAccountCachedBalance(state) {
-  const cachedBalances = state.metamask.cachedBalances[state.metamask.network];
+  const cachedBalances = getMetaMaskCachedBalances(state);
   const selectedAddress = getSelectedAddress(state);
 
   return cachedBalances && cachedBalances[selectedAddress];

ui/app/selectors/tests/permissions.test.js

@@ -1,4 +1,5 @@
 import assert from 'assert';
+import { KOVAN_CHAIN_ID } from '../../../../shared/constants/network';
 import {
   getConnectedDomainsForSelectedAddress,
   getOrderedConnectedAccountsForActiveTab,
@@ -163,6 +164,9 @@ describe('selectors', function () {
         url: 'https://remix.ethereum.org/',
       },
       metamask: {
+        provider: {
+          chainId: KOVAN_CHAIN_ID,
+        },
         accounts: {
           0x7250739de134d33ec7ab1ee592711e15098c9d2d: {
             address: '0x7250739de134d33ec7ab1ee592711e15098c9d2d',