add unit tests

MetaMask · Nov 13, 2024 · 94a24d2 · 94a24d2
1 parent f8cc287
commit 94a24d2
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 0 deletions.
diff --git a/packages/assets-controllers/src/TokensController.test.ts b/packages/assets-controllers/src/TokensController.test.ts
@@ -17,8 +17,10 @@ import type { InternalAccount } from '@metamask/keyring-api';
 import type {
   NetworkClientConfiguration,
   NetworkClientId,
+  NetworkState,
 } from '@metamask/network-controller';
 import { getDefaultNetworkControllerState } from '@metamask/network-controller';
+import type { Patch } from 'immer';
 import nock from 'nock';
 import * as sinon from 'sinon';
 import { v1 as uuidV1 } from 'uuid';
@@ -2447,6 +2449,64 @@ describe('TokensController', () => {
     });
   });
 
+  describe('when NetworkController:stateChange is published', () => {
+    it('removes tokens for removed networks', async () => {
+      const initialState = {
+        allTokens: {
+          '0x1': {
+            '0x134': [
+              {
+                address: '0x01',
+                symbol: 'TKN1',
+                decimals: 18,
+                aggregators: [],
+                name: 'Token 1',
+              },
+            ],
+          },
+          '0x5': {
+            // goerli
+            '0x456': [
+              {
+                address: '0x02',
+                symbol: 'TKN2',
+                decimals: 18,
+                aggregators: [],
+                name: 'Token 2',
+              },
+            ],
+          },
+        },
+        tokens: [],
+        ignoredTokens: [],
+        detectedTokens: [],
+        allIgnoredTokens: {},
+        allDetectedTokens: {},
+      };
+
+      await withController(
+        { options: { state: initialState } },
+        async ({ controller, triggerNetworkStateChange }) => {
+          // Verify initial state
+          expect(controller.state).toStrictEqual(initialState);
+
+          // Simulate removing goerli
+          triggerNetworkStateChange({} as NetworkState, [
+            {
+              op: 'remove',
+              path: ['networkConfigurationsByChainId', '0x5'],
+            },
+          ]);
+
+          // Verify tokens were removed on goerli
+          expect(controller.state.allTokens).toStrictEqual({
+            '0x1': initialState.allTokens['0x1'],
+          });
+        },
+      );
+    });
+  });
+
   describe('resetState', () => {
     it('resets the state to default state', async () => {
       const initialState: TokensControllerState = {
@@ -2544,6 +2604,10 @@ type WithControllerCallback<ReturnValue> = ({
   messenger: UnrestrictedMessenger;
   approvalController: ApprovalController;
   triggerSelectedAccountChange: (internalAccount: InternalAccount) => void;
+  triggerNetworkStateChange: (
+    networkState: NetworkState,
+    patches: Patch[],
+  ) => void;
   getAccountHandler: jest.Mock;
   getSelectedAccountHandler: jest.Mock;
 }) => Promise<ReturnValue> | ReturnValue;
@@ -2675,12 +2739,20 @@ async function withController<ReturnValue>(
     getNetworkClientById,
   );
 
+  const triggerNetworkStateChange = (
+    networkState: NetworkState,
+    patches: Patch[],
+  ) => {
+    messenger.publish('NetworkController:stateChange', networkState, patches);
+  };
+
   return await fn({
     controller,
     changeNetwork,
     messenger,
     approvalController,
     triggerSelectedAccountChange,
+    triggerNetworkStateChange,
     getAccountHandler,
     getSelectedAccountHandler,
   });

diff --git a/packages/assets-controllers/src/multicall.test.ts b/packages/assets-controllers/src/multicall.test.ts
@@ -67,6 +67,35 @@ describe('multicall', () => {
       ]);
     });
 
+    it('should handle the multicall contract returning false for success', async () => {
+      // Mock an unsuccessful multicall
+      jest
+        .spyOn(provider, 'call')
+        .mockResolvedValue(
+          defaultAbiCoder.encode(
+            ['tuple(bool,bytes)[]'],
+            [
+              calls.map((_, i) => [
+                false,
+                defaultAbiCoder.encode(['uint256'], [i + 1]),
+              ]),
+            ],
+          ),
+        );
+
+      const results = await multicallOrFallback(calls, '0x1', provider);
+      expect(results).toMatchObject([
+        {
+          success: false,
+          value: undefined,
+        },
+        {
+          success: false,
+          value: undefined,
+        },
+      ]);
+    });
+
     it('should fallback to parallel calls on unsupported chains', async () => {
       // Mock return values for each call
       let timesCalled = 0;