feat: upgrade assets controllers to v44

app/components/UI/AccountFromToInfoCard/AccountFromToInfoCard.test.tsx

@@ -38,9 +38,7 @@ const mockInitialState: DeepPartial<RootState> = {
         },
       },
       TokenBalancesController: {
-        contractBalances: {
-          '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': '0x5',
-        },
+        tokenBalances: { },
       },
       AccountsController: MOCK_ACCOUNTS_CONTROLLER_STATE,
     },

app/components/UI/AccountInfoCard/index.test.tsx

@@ -67,7 +67,7 @@ const mockInitialState: DeepPartial<RootState> = {
         }),
       },
       TokenBalancesController: {
-        contractBalances: {},
+        tokenBalances: {},
       },
     },
   },

app/components/UI/Tokens/TokenList/PortfolioBalance/index.test.tsx

@@ -1,6 +1,5 @@
 import React from 'react';
 import { fireEvent } from '@testing-library/react-native';
-import { BN } from 'ethereumjs-util';
 import renderWithProvider from '../../../../../util/test/renderWithProvider';
 import { backgroundState } from '../../../../../util/test/initial-root-state';
 import AppConstants from '../../../../../../app/core/AppConstants';
@@ -76,11 +75,7 @@ const initialState = {
         },
       },
       TokenBalancesController: {
-        contractBalances: {
-          '0x00': new BN(2),
-          '0x01': new BN(2),
-          '0x02': new BN(0),
-        },
+        tokenBalances: { },
       },
     },
   },

app/components/UI/Tokens/index.test.tsx

@@ -56,10 +56,22 @@ jest.mock('../../../core/Engine', () => ({
   },
 }));
 
+const selectedAddress = '0x123';
+
 const initialState = {
   engine: {
     backgroundState: {
       ...backgroundState,
+      AccountsController: {
+        internalAccounts: {
+          selectedAccount: '1',
+          accounts: {
+            '1': {
+              address: selectedAddress,
+            },
+          },
+        },
+      },
       TokensController: {
         tokens: [
           {
@@ -109,11 +121,15 @@ const initialState = {
         },
       },
       TokenBalancesController: {
-        contractBalances: {
-          '0x00': new BN(2),
-          '0x01': new BN(2),
-          '0x02': new BN(0),
-        },
+        tokenBalances:{
+          [selectedAddress]: {
+            '0x1': {
+              '0x00': new BN(2),
+              '0x01': new BN(2),
+              '0x02': new BN(0),
+            }
+          }
+        }
       },
     },
   },
@@ -263,9 +279,23 @@ describe('Tokens', () => {
               },
             },
           },
+          AccountsController: {
+            internalAccounts: {
+              selectedAccount: '1',
+              accounts: {
+                '1': {
+                  address: selectedAddress,
+                },
+              },
+            },
+          },
           TokenBalancesController: {
-            contractBalances: {
-              '0x02': new BN(1),
+            tokenBalances: {
+              [selectedAddress]: {
+                '0x1': {
+                  '0x02': new BN(1),
+                },
+              },
             },
           },
         },

app/components/Views/DetectedTokens/components/Token.tsx

@@ -127,7 +127,7 @@ const Token = ({ token, selected, toggleSelected }: Props) => {
     (tokenExchangeRates as Record<Hex, MarketDataDetails>)?.[address as Hex] ??
     null;
   const tokenBalance = renderFromTokenMinimalUnit(
-    tokenBalances[address],
+    tokenBalances[address as Hex],
     decimals,
   );
   const tokenBalanceWithSymbol = `${

app/components/Views/SDK/SDKSessionsManager/SDKSessionAccountListItem.test.tsx

@@ -23,9 +23,7 @@ const mockInitialState: DeepPartial<RootState> = {
         },
       },
       TokenBalancesController: {
-        contractBalances: {
-          '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': '0x5',
-        },
+        tokenBalances: { },
       },
       PreferencesController: {
         selectedAddress: '0xe64dD0AB5ad7e8C5F2bf6Ce75C34e187af8b920A',

app/components/Views/confirmations/Send/index.test.tsx

@@ -58,7 +58,7 @@ const initialState: DeepPartial<RootState> = {
         addressBook: {},
       },
       TokenBalancesController: {
-        contractBalances: {},
+        tokenBalances: {},
       },
       TokenListController: {
         tokenList: { '0x1': {} },

app/components/Views/confirmations/components/TransactionReview/TransactionReviewEIP1559Update/index.test.tsx

@@ -38,9 +38,7 @@ const initialState: DeepPartial<RootState> = {
         },
       },
       TokenBalancesController: {
-        contractBalances: {
-          '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': '0x5',
-        },
+        tokenBalances: { },
       },
       CurrencyRateController: {
         currentCurrency: 'usd',

app/components/hooks/AssetPolling/AssetPollingProvider.test.tsx

@@ -0,0 +1,27 @@
+import React from 'react';
+import { render } from '@testing-library/react-native';
+
+import { AssetPollingProvider } from './AssetPollingProvider';
+
+jest.mock('./useCurrencyRatePolling', () => jest.fn());
+jest.mock('./useTokenRatesPolling', () => jest.fn());
+jest.mock('./useTokenDetectionPolling', () => jest.fn());
+jest.mock('./useTokenListPolling', () => jest.fn());
+jest.mock('./useTokenBalancesPolling', () => jest.fn());
+
+describe('AssetPollingProvider', () => {
+  it('should call all polling hooks', () => {
+
+    render(
+      <AssetPollingProvider>
+        <div></div>
+      </AssetPollingProvider>
+    );
+
+    expect(jest.requireMock('./useCurrencyRatePolling')).toHaveBeenCalled();
+    expect(jest.requireMock('./useTokenRatesPolling')).toHaveBeenCalled();
+    expect(jest.requireMock('./useTokenDetectionPolling')).toHaveBeenCalled();
+    expect(jest.requireMock('./useTokenListPolling')).toHaveBeenCalled();
+    expect(jest.requireMock('./useTokenBalancesPolling')).toHaveBeenCalled();
+  });
+});

app/components/hooks/AssetPolling/AssetPollingProvider.tsx

@@ -3,6 +3,7 @@ import useCurrencyRatePolling from './useCurrencyRatePolling';
 import useTokenRatesPolling from './useTokenRatesPolling';
 import useTokenDetectionPolling from './useTokenDetectionPolling';
 import useTokenListPolling from './useTokenListPolling';
+import useTokenBalancesPolling from './useTokenBalancesPolling';
 
 // This provider is a step towards making controller polling fully UI based.
 // Eventually, individual UI components will call the use*Polling hooks to
@@ -12,6 +13,7 @@ export const AssetPollingProvider = ({ children }: { children: ReactNode }) => {
   useTokenRatesPolling();
   useTokenDetectionPolling();
   useTokenListPolling();
+  useTokenBalancesPolling();
 
   return <>{children}</>;
 };

app/components/hooks/AssetPolling/useTokenBalancesPolling.test.ts

@@ -0,0 +1,58 @@
+import { renderHookWithProvider } from '../../../util/test/renderWithProvider';
+import Engine from '../../../core/Engine';
+import useTokenBalancesPolling from './useTokenBalancesPolling';
+
+jest.mock('../../../core/Engine', () => ({
+  context: {
+    TokenBalancesController: {
+      startPolling: jest.fn(),
+      stopPollingByPollingToken: jest.fn(),
+    },
+  },
+}));
+
+describe('useTokenBalancesPolling', () => {
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+  });
+
+  const selectedChainId = '0x1' as const;
+  const state = {
+    engine: {
+      backgroundState: {
+        TokenBalancesController: {
+          tokenBalances: {},
+        },
+        NetworkController: {
+          selectedNetworkClientId: 'selectedNetworkClientId',
+          networkConfigurationsByChainId: {
+            [selectedChainId]: {
+              chainId: selectedChainId,
+              rpcEndpoints: [{
+                networkClientId: 'selectedNetworkClientId',
+              }]
+            },
+            '0x89': {},
+          },
+        },
+      },
+    },
+  };
+
+  it('Should poll by selected chain id, and stop polling on dismount', async () => {
+
+    const { unmount } = renderHookWithProvider(() => useTokenBalancesPolling(), {state});
+
+    const mockedTokenBalancesController = jest.mocked(Engine.context.TokenBalancesController);
+
+    expect(mockedTokenBalancesController.startPolling).toHaveBeenCalledTimes(1);
+    expect(
+      mockedTokenBalancesController.startPolling
+    ).toHaveBeenCalledWith({chainId: selectedChainId});
+
+    expect(mockedTokenBalancesController.stopPollingByPollingToken).toHaveBeenCalledTimes(0);
+    unmount();
+    expect(mockedTokenBalancesController.stopPollingByPollingToken).toHaveBeenCalledTimes(1);
+  });
+});

app/components/hooks/AssetPolling/useTokenBalancesPolling.ts

@@ -0,0 +1,37 @@
+import { useSelector } from 'react-redux';
+import usePolling from '../usePolling';
+import Engine from '../../../core/Engine';
+import { selectChainId, selectNetworkConfigurations } from '../../../selectors/networkController';
+import { Hex } from '@metamask/utils';
+import { isPortfolioViewEnabled } from '../../../util/networks';
+import { selectAllTokenBalances } from '../../../selectors/tokenBalancesController';
+
+const useTokenBalancesPolling = ({ chainIds }: { chainIds?: Hex[] } = {}) => {
+
+  // Selectors to determine polling input
+  const networkConfigurations = useSelector(selectNetworkConfigurations);
+  const currentChainId = useSelector(selectChainId);
+
+  // Selectors returning state updated by the polling
+  const tokenBalances = useSelector(selectAllTokenBalances);
+
+  const chainIdsToPoll = isPortfolioViewEnabled
+    ? (chainIds ?? Object.keys(networkConfigurations))
+    : [currentChainId];
+
+  const { TokenBalancesController } = Engine.context;
+
+  usePolling({
+    startPolling:
+      TokenBalancesController.startPolling.bind(TokenBalancesController),
+    stopPollingByPollingToken:
+      TokenBalancesController.stopPollingByPollingToken.bind(TokenBalancesController),
+    input: chainIdsToPoll.map((chainId) => ({chainId: chainId as Hex})),
+  });
+
+  return {
+    tokenBalances
+  };
+};
+
+export default useTokenBalancesPolling;

app/components/hooks/useAddressBalance/useAddressBalance.test.tsx

@@ -34,8 +34,12 @@ const mockInitialState = {
         },
       },
       TokenBalancesController: {
-        contractBalances: {
-          '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': '0x5',
+        tokenBalances: {
+          [MOCK_ADDRESS_1]: {
+            '0x1': {
+              '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': '0x5',
+            },
+          },
         },
       },
       PreferencesController: {

app/components/hooks/useTokenBalancesController/useTokenBalancesController.test.tsx

@@ -8,14 +8,31 @@ import { BN } from 'ethereumjs-util';
 import { cloneDeep } from 'lodash';
 import { backgroundState } from '../../../util/test/initial-root-state';
 
+const accountAddress = '0x123';
+const chainId = '0x1';
+
 // initial state for the test store
 const mockInitialState = {
   engine: {
     backgroundState: {
       ...backgroundState,
+      AccountsController: {
+        internalAccounts: {
+          selectedAccount: '1',
+          accounts: {
+            '1': {
+              address: accountAddress,
+            },
+          },
+        },
+      },
       TokenBalancesController: {
-        contractBalances: {
-          '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': new BN(0x2a),
+        tokenBalances: {
+          [accountAddress]: {
+            [chainId]: {
+              '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': new BN(0x2a),
+            }
+          }
         },
       },
     },
@@ -35,10 +52,20 @@ const testBalancesReducer = (state: any, action: any) => {
           ...state.engine.backgroundState,
           TokenBalancesController: {
             ...state.engine.backgroundState.TokenBalancesController,
-            contractBalances: {
+            tokenBalances: {
               ...state.engine.backgroundState.TokenBalancesController
-                .contractBalances,
+                .tokenBalances,
               ...action.value,
+              [accountAddress]: {
+                ...state.engine.backgroundState.TokenBalancesController
+                  .tokenBalances[accountAddress],
+                ...action.value[accountAddress],
+                [chainId]: {
+                  ...state.engine.backgroundState.TokenBalancesController
+                    .tokenBalances[accountAddress][chainId],
+                  ...action.value[accountAddress][chainId],
+                }
+              },
             },
           },
         },
@@ -104,7 +131,11 @@ describe('useTokenBalancesController()', () => {
       testStore.dispatch({
         type: 'add-balances',
         value: {
-          '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b96': new BN(0x539),
+          [accountAddress]: {
+            [chainId]: {
+              '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b96': new BN(0x539),
+            }
+          }
         },
       });
     });
@@ -122,7 +153,11 @@ describe('useTokenBalancesController()', () => {
       testStore.dispatch({
         type: 'add-balances',
         value: {
-          '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': new BN(0x2a),
+          [accountAddress]: {
+            [chainId]: {
+              '0x326836cc6cd09B5aa59B81A7F72F25FcC0136b95': new BN(0x2a),
+            }
+          }
         },
       });
     });