From 37edcfc91c5abe475922f0fe8307eabc172f6f73 Mon Sep 17 00:00:00 2001 From: Salah-Eddine Saakoun Date: Thu, 31 Oct 2024 11:35:25 +0100 Subject: [PATCH 1/4] fix: add missing resetState methods --- .../src/NftController.test.ts | 41 ++++++++++ .../assets-controllers/src/NftController.ts | 11 +++ .../src/TokenBalancesController.test.ts | 29 +++++++ .../src/TokenBalancesController.ts | 9 ++ .../src/TokenRatesController.test.ts | 51 ++++++++++++ .../src/TokenRatesController.ts | 9 ++ .../src/TokensController.test.ts | 82 +++++++++++++++++++ .../src/TokensController.ts | 14 ++++ 8 files changed, 246 insertions(+) diff --git a/packages/assets-controllers/src/NftController.test.ts b/packages/assets-controllers/src/NftController.test.ts index 55adcc3a3e..2fd003dac4 100644 --- a/packages/assets-controllers/src/NftController.test.ts +++ b/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: [], + }); + }); + }); }); diff --git a/packages/assets-controllers/src/NftController.ts b/packages/assets-controllers/src/NftController.ts index 861f657df5..c8023b60b8 100644 --- a/packages/assets-controllers/src/NftController.ts +++ b/packages/assets-controllers/src/NftController.ts @@ -2060,6 +2060,17 @@ export class NftController extends BaseController< }); } } + + /** + * Reset the controller state to the initial state. + */ + resetState() { + this.update((state) => { + state.allNftContracts = {}; + state.allNfts = {}; + state.ignoredNfts = []; + }); + } } export default NftController; diff --git a/packages/assets-controllers/src/TokenBalancesController.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts index e57d13c77f..b2099d05c1 100644 --- a/packages/assets-controllers/src/TokenBalancesController.test.ts +++ b/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: {}, + }); + }); + }); }); diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index e80db70d5f..345f5e0372 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/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((state) => { + state.contractBalances = {}; + }); + } } export default TokenBalancesController; diff --git a/packages/assets-controllers/src/TokenRatesController.test.ts b/packages/assets-controllers/src/TokenRatesController.test.ts index 5387713024..20b040f75e 100644 --- a/packages/assets-controllers/src/TokenRatesController.test.ts +++ b/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. diff --git a/packages/assets-controllers/src/TokenRatesController.ts b/packages/assets-controllers/src/TokenRatesController.ts index 52d51b0d3f..789d50a496 100644 --- a/packages/assets-controllers/src/TokenRatesController.ts +++ b/packages/assets-controllers/src/TokenRatesController.ts @@ -756,6 +756,15 @@ export class TokenRatesController extends StaticIntervalPollingController { + state.marketData = {}; + }); + } } export default TokenRatesController; diff --git a/packages/assets-controllers/src/TokensController.test.ts b/packages/assets-controllers/src/TokensController.test.ts index a0b645d1a9..eb12ef587d 100644 --- a/packages/assets-controllers/src/TokensController.test.ts +++ b/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 = ({ diff --git a/packages/assets-controllers/src/TokensController.ts b/packages/assets-controllers/src/TokensController.ts index 0f74f8b878..0b394c4055 100644 --- a/packages/assets-controllers/src/TokensController.ts +++ b/packages/assets-controllers/src/TokensController.ts @@ -1077,6 +1077,20 @@ export class TokensController extends BaseController< ); return account?.address || ''; } + + /** + * Reset the controller state to the initial state. + */ + resetState() { + this.update((state) => { + state.tokens = []; + state.ignoredTokens = []; + state.detectedTokens = []; + state.allTokens = {}; + state.allIgnoredTokens = {}; + state.allDetectedTokens = {}; + }); + } } export default TokensController; From f4a7e841aae93bb69ce1e51ab0877c83b13b61c4 Mon Sep 17 00:00:00 2001 From: Salah-Eddine Saakoun Date: Thu, 31 Oct 2024 11:52:04 +0100 Subject: [PATCH 2/4] fix: add CHANGELOG entry --- packages/assets-controllers/CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md index 5c0b9fc657..4467d272ac 100644 --- a/packages/assets-controllers/CHANGELOG.md +++ b/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 initial state ([#4880](https://github.com/MetaMask/core/pull/4880)) + ## [41.0.0] ### Changed From d0b8a4edf791c362dabd34da3f8454b7b8b051d7 Mon Sep 17 00:00:00 2001 From: Salah-Eddine Saakoun Date: Thu, 31 Oct 2024 20:38:48 +0100 Subject: [PATCH 3/4] fix: apply review feedback --- packages/assets-controllers/CHANGELOG.md | 2 +- packages/assets-controllers/src/NftController.ts | 6 ++---- .../assets-controllers/src/TokenBalancesController.ts | 4 ++-- packages/assets-controllers/src/TokenRatesController.ts | 4 ++-- packages/assets-controllers/src/TokensController.ts | 9 ++------- 5 files changed, 9 insertions(+), 16 deletions(-) diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md index 4467d272ac..7e988b889c 100644 --- a/packages/assets-controllers/CHANGELOG.md +++ b/packages/assets-controllers/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add `resetState` method to `NftController`, `TokensController`, `TokenBalancesController` and `TokenRatesController` to reset the controller's state back to their initial state ([#4880](https://github.com/MetaMask/core/pull/4880)) +- 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] diff --git a/packages/assets-controllers/src/NftController.ts b/packages/assets-controllers/src/NftController.ts index c8023b60b8..3d0ffc9f8f 100644 --- a/packages/assets-controllers/src/NftController.ts +++ b/packages/assets-controllers/src/NftController.ts @@ -2065,10 +2065,8 @@ export class NftController extends BaseController< * Reset the controller state to the initial state. */ resetState() { - this.update((state) => { - state.allNftContracts = {}; - state.allNfts = {}; - state.ignoredNfts = []; + this.update(() => { + return getDefaultNftControllerState(); }); } } diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 345f5e0372..40b74f786f 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -225,8 +225,8 @@ export class TokenBalancesController extends BaseController< * Reset the controller state to the initial state. */ resetState() { - this.update((state) => { - state.contractBalances = {}; + this.update(() => { + return getDefaultTokenBalancesState(); }); } } diff --git a/packages/assets-controllers/src/TokenRatesController.ts b/packages/assets-controllers/src/TokenRatesController.ts index 789d50a496..30aace943a 100644 --- a/packages/assets-controllers/src/TokenRatesController.ts +++ b/packages/assets-controllers/src/TokenRatesController.ts @@ -761,8 +761,8 @@ export class TokenRatesController extends StaticIntervalPollingController { - state.marketData = {}; + this.update(() => { + return getDefaultTokenRatesControllerState(); }); } } diff --git a/packages/assets-controllers/src/TokensController.ts b/packages/assets-controllers/src/TokensController.ts index 0b394c4055..2444302fbe 100644 --- a/packages/assets-controllers/src/TokensController.ts +++ b/packages/assets-controllers/src/TokensController.ts @@ -1082,13 +1082,8 @@ export class TokensController extends BaseController< * Reset the controller state to the initial state. */ resetState() { - this.update((state) => { - state.tokens = []; - state.ignoredTokens = []; - state.detectedTokens = []; - state.allTokens = {}; - state.allIgnoredTokens = {}; - state.allDetectedTokens = {}; + this.update(() => { + return getDefaultTokensState(); }); } } From dd732dd20532f89bf125e6142a84f06543145e39 Mon Sep 17 00:00:00 2001 From: Salah-Eddine Saakoun Date: Thu, 31 Oct 2024 21:00:10 +0100 Subject: [PATCH 4/4] fix: use default instead initial --- packages/assets-controllers/src/NftController.ts | 2 +- packages/assets-controllers/src/TokenBalancesController.ts | 2 +- packages/assets-controllers/src/TokenRatesController.ts | 2 +- packages/assets-controllers/src/TokensController.ts | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/assets-controllers/src/NftController.ts b/packages/assets-controllers/src/NftController.ts index 3d0ffc9f8f..63c60b53aa 100644 --- a/packages/assets-controllers/src/NftController.ts +++ b/packages/assets-controllers/src/NftController.ts @@ -2062,7 +2062,7 @@ export class NftController extends BaseController< } /** - * Reset the controller state to the initial state. + * Reset the controller state to the default state. */ resetState() { this.update(() => { diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 40b74f786f..0c0cb505ef 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -222,7 +222,7 @@ export class TokenBalancesController extends BaseController< } /** - * Reset the controller state to the initial state. + * Reset the controller state to the default state. */ resetState() { this.update(() => { diff --git a/packages/assets-controllers/src/TokenRatesController.ts b/packages/assets-controllers/src/TokenRatesController.ts index 30aace943a..ec2ade5e0b 100644 --- a/packages/assets-controllers/src/TokenRatesController.ts +++ b/packages/assets-controllers/src/TokenRatesController.ts @@ -758,7 +758,7 @@ export class TokenRatesController extends StaticIntervalPollingController { diff --git a/packages/assets-controllers/src/TokensController.ts b/packages/assets-controllers/src/TokensController.ts index 2444302fbe..fe680cdd0d 100644 --- a/packages/assets-controllers/src/TokensController.ts +++ b/packages/assets-controllers/src/TokensController.ts @@ -1079,7 +1079,7 @@ export class TokensController extends BaseController< } /** - * Reset the controller state to the initial state. + * Reset the controller state to the default state. */ resetState() { this.update(() => {