fix: add missing resetState methods to 4 assets-controllers

packages/assets-controllers/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add `resetState` method to `NftController`, `TokensController`, `TokenBalancesController` and `TokenRatesController` to reset the controller's state back to their default state ([#4880](https://github.com/MetaMask/core/pull/4880))
+
 ## [41.0.0]
 
 ### Changed

packages/assets-controllers/src/NftController.test.ts

@@ -4944,4 +4944,45 @@ describe('NftController', () => {
       });
     });
   });
+
+  describe('resetState', () => {
+    it('resets the state to default state', () => {
+      const initialState: NftControllerState = {
+        allNftContracts: {
+          [OWNER_ACCOUNT.address]: { [ChainId.mainnet]: [] },
+        },
+        allNfts: {
+          [OWNER_ACCOUNT.address]: { [ChainId.mainnet]: [] },
+        },
+        ignoredNfts: [
+          {
+            address: ERC1155_NFT_ADDRESS,
+            name: null,
+            description: null,
+            image: null,
+            tokenId: ERC1155_NFT_ID,
+            standard: ERC1155,
+            favorite: false,
+            isCurrentlyOwned: true,
+            tokenURI: 'ipfs://*',
+          },
+        ],
+      };
+      const { nftController } = setupController({
+        options: {
+          state: initialState,
+        },
+      });
+
+      expect(nftController.state).toStrictEqual(initialState);
+
+      nftController.resetState();
+
+      expect(nftController.state).toStrictEqual({
+        allNftContracts: {},
+        allNfts: {},
+        ignoredNfts: [],
+      });
+    });
+  });
 });

packages/assets-controllers/src/NftController.ts

@@ -2060,6 +2060,15 @@ export class NftController extends BaseController<
       });
     }
   }
+
+  /**
+   * Reset the controller state to the initial state.
+   */
+  resetState() {
+    this.update(() => {
+      return getDefaultNftControllerState();
+    });
+  }
 }
 
 export default NftController;

packages/assets-controllers/src/TokenBalancesController.test.ts

@@ -11,6 +11,7 @@ import type {
   TokenBalancesControllerActions,
   TokenBalancesControllerEvents,
   TokenBalancesControllerMessenger,
+  TokenBalancesControllerState,
 } from './TokenBalancesController';
 import { TokenBalancesController } from './TokenBalancesController';
 import type { Token } from './TokenRatesController';
@@ -457,4 +458,32 @@ describe('TokenBalancesController', () => {
       '0x02': toHex(new BN(1)),
     });
   });
+
+  describe('resetState', () => {
+    it('resets the state to default state', () => {
+      const initialState: TokenBalancesControllerState = {
+        contractBalances: {
+          '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0': toHex(new BN(1)),
+        },
+      };
+
+      const { controller } = setupController({
+        config: {
+          state: initialState,
+          disabled: true,
+        },
+        mock: {
+          selectedAccount: createMockInternalAccount({ address: '0x1234' }),
+        },
+      });
+
+      expect(controller.state).toStrictEqual(initialState);
+
+      controller.resetState();
+
+      expect(controller.state).toStrictEqual({
+        contractBalances: {},
+      });
+    });
+  });
 });

packages/assets-controllers/src/TokenBalancesController.ts

@@ -220,6 +220,15 @@ export class TokenBalancesController extends BaseController<
       state.contractBalances = newContractBalances;
     });
   }
+
+  /**
+   * Reset the controller state to the initial state.
+   */
+  resetState() {
+    this.update(() => {
+      return getDefaultTokenBalancesState();
+    });
+  }
 }
 
 export default TokenBalancesController;

packages/assets-controllers/src/TokenRatesController.test.ts

@@ -38,6 +38,7 @@ import type {
   AllowedEvents,
   Token,
   TokenRatesControllerMessenger,
+  TokenRatesControllerState,
 } from './TokenRatesController';
 import { getDefaultTokensState } from './TokensController';
 import type { TokensControllerState } from './TokensController';
@@ -2341,6 +2342,56 @@ describe('TokenRatesController', () => {
       });
     });
   });
+
+  describe('resetState', () => {
+    it('resets the state to default state', async () => {
+      const initialState: TokenRatesControllerState = {
+        marketData: {
+          [ChainId.mainnet]: {
+            '0x02': {
+              currency: 'ETH',
+              priceChange1d: 0,
+              pricePercentChange1d: 0,
+              tokenAddress: '0x02',
+              allTimeHigh: 4000,
+              allTimeLow: 900,
+              circulatingSupply: 2000,
+              dilutedMarketCap: 100,
+              high1d: 200,
+              low1d: 100,
+              marketCap: 1000,
+              marketCapPercentChange1d: 100,
+              price: 0.001,
+              pricePercentChange14d: 100,
+              pricePercentChange1h: 1,
+              pricePercentChange1y: 200,
+              pricePercentChange200d: 300,
+              pricePercentChange30d: 200,
+              pricePercentChange7d: 100,
+              totalVolume: 100,
+            },
+          },
+        },
+      };
+
+      await withController(
+        {
+          options: {
+            state: initialState,
+          },
+        },
+        ({ controller }) => {
+          expect(controller.state).toStrictEqual(initialState);
+
+          controller.resetState();
+
+          expect(controller.state).toStrictEqual({
+            marketData: {},
+          });
+        },
+      );
+    });
+  });
 });
 /**
  * A callback for the `withController` helper function.

packages/assets-controllers/src/TokenRatesController.ts

@@ -756,6 +756,15 @@ export class TokenRatesController extends StaticIntervalPollingController<TokenR
 
     return updatedContractExchangeRates;
   }
+
+  /**
+   * Reset the controller state to the initial state.
+   */
+  resetState() {
+    this.update(() => {
+      return getDefaultTokenRatesControllerState();
+    });
+  }
 }
 
 export default TokenRatesController;

packages/assets-controllers/src/TokensController.test.ts

@@ -2308,6 +2308,88 @@ describe('TokensController', () => {
       });
     });
   });
+
+  describe('resetState', () => {
+    it('resets the state to default state', async () => {
+      const initialState: TokensControllerState = {
+        detectedTokens: [
+          {
+            address: '0x01',
+            symbol: 'barA',
+            decimals: 2,
+            aggregators: [],
+            image: undefined,
+            name: undefined,
+          },
+        ],
+        tokens: [
+          {
+            address: '0x02',
+            symbol: 'barB',
+            decimals: 2,
+            aggregators: [],
+            image: undefined,
+            name: undefined,
+          },
+        ],
+        allTokens: {
+          [ChainId.mainnet]: {
+            '0x0001': [
+              {
+                address: '0x03',
+                symbol: 'barC',
+                decimals: 2,
+                aggregators: [],
+                image: undefined,
+                name: undefined,
+              },
+            ],
+          },
+        },
+        ignoredTokens: ['0x03'],
+        allIgnoredTokens: {
+          [ChainId.mainnet]: {
+            '0x0001': ['0x03'],
+          },
+        },
+        allDetectedTokens: {
+          [ChainId.mainnet]: {
+            '0x0001': [
+              {
+                address: '0x01',
+                symbol: 'barA',
+                decimals: 2,
+                aggregators: [],
+                image: undefined,
+                name: undefined,
+              },
+            ],
+          },
+        },
+      };
+      await withController(
+        {
+          options: {
+            state: initialState,
+          },
+        },
+        ({ controller }) => {
+          expect(controller.state).toStrictEqual(initialState);
+
+          controller.resetState();
+
+          expect(controller.state).toStrictEqual({
+            tokens: [],
+            ignoredTokens: [],
+            detectedTokens: [],
+            allTokens: {},
+            allIgnoredTokens: {},
+            allDetectedTokens: {},
+          });
+        },
+      );
+    });
+  });
 });
 
 type WithControllerCallback<ReturnValue> = ({

packages/assets-controllers/src/TokensController.ts

@@ -1077,6 +1077,15 @@ export class TokensController extends BaseController<
     );
     return account?.address || '';
   }
+
+  /**
+   * Reset the controller state to the initial state.
+   */
+  resetState() {
+    this.update(() => {
+      return getDefaultTokensState();
+    });
+  }
 }
 
 export default TokensController;