feat: clear state on SubscriptionsController, ClaimsController and ShieldController

packages/claims-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Added new public method, `clearState` to clear/reset the claims controller state. ([#7780](https://github.com/MetaMask/core/pull/7780))
+
 ### Changed
 
 - Bump `@metamask/controller-utils` from `^11.17.0` to `^11.18.0` ([#7583](https://github.com/MetaMask/core/pull/7583))

packages/claims-controller/src/ClaimsController.test.ts

@@ -1,6 +1,9 @@
 import { toHex } from '@metamask/controller-utils';
 
-import { ClaimsController } from './ClaimsController';
+import {
+  ClaimsController,
+  getDefaultClaimsControllerState,
+} from './ClaimsController';
 import { ClaimsControllerErrorMessages, ClaimStatusEnum } from './constants';
 import type {
   Claim,
@@ -18,6 +21,35 @@ const mockKeyringControllerSignPersonalMessage = jest.fn();
 const mockClaimsServiceGetClaims = jest.fn();
 const mockClaimsServiceFetchClaimsConfigurations = jest.fn();
 
+const MOCK_CLAIM_1: Claim = {
+  id: 'mock-claim-1',
+  shortId: 'mock-claim-1',
+  status: ClaimStatusEnum.CREATED,
+  createdAt: '2021-01-01',
+  updatedAt: '2021-01-01',
+  chainId: '0x1',
+  email: 'test@test.com',
+  impactedWalletAddress: '0x123',
+  impactedTxHash: '0x123',
+  reimbursementWalletAddress: '0x456',
+  description: 'test description',
+  signature: '0xdeadbeef',
+};
+const MOCK_CLAIM_2: Claim = {
+  id: 'mock-claim-2',
+  shortId: 'mock-claim-2',
+  status: ClaimStatusEnum.CREATED,
+  createdAt: '2021-01-01',
+  updatedAt: '2021-01-01',
+  chainId: '0x1',
+  email: 'test2@test.com',
+  impactedWalletAddress: '0x789',
+  impactedTxHash: '0x789',
+  reimbursementWalletAddress: '0x012',
+  description: 'test description 2',
+  signature: '0xdeadbeef',
+};
+
 /**
  * Builds a controller based on the given options and calls the given function with that controller.
  *
@@ -210,35 +242,6 @@ describe('ClaimsController', () => {
   });
 
   describe('getClaims', () => {
-    const MOCK_CLAIM_1: Claim = {
-      id: 'mock-claim-1',
-      shortId: 'mock-claim-1',
-      status: ClaimStatusEnum.CREATED,
-      createdAt: '2021-01-01',
-      updatedAt: '2021-01-01',
-      chainId: '0x1',
-      email: 'test@test.com',
-      impactedWalletAddress: '0x123',
-      impactedTxHash: '0x123',
-      reimbursementWalletAddress: '0x456',
-      description: 'test description',
-      signature: '0xdeadbeef',
-    };
-    const MOCK_CLAIM_2: Claim = {
-      id: 'mock-claim-2',
-      shortId: 'mock-claim-2',
-      status: ClaimStatusEnum.CREATED,
-      createdAt: '2021-01-01',
-      updatedAt: '2021-01-01',
-      chainId: '0x1',
-      email: 'test2@test.com',
-      impactedWalletAddress: '0x789',
-      impactedTxHash: '0x789',
-      reimbursementWalletAddress: '0x012',
-      description: 'test description 2',
-      signature: '0xdeadbeef',
-    };
-
     it('should be able to get the list of claims', async () => {
       await withController(async ({ controller }) => {
         mockClaimsServiceGetClaims.mockResolvedValueOnce([
@@ -381,4 +384,32 @@ describe('ClaimsController', () => {
       );
     });
   });
+
+  describe('clearState', () => {
+    it('should reset state to default values', async () => {
+      await withController(
+        {
+          state: {
+            claims: [MOCK_CLAIM_1, MOCK_CLAIM_2],
+            drafts: [MOCK_CLAIM_1, MOCK_CLAIM_2].map((claim) => ({
+              draftId: claim.id,
+              ...claim,
+            })),
+          },
+        },
+        async ({ controller }) => {
+          expect(controller.state.claims).toHaveLength(2);
+          expect(controller.state.drafts).toHaveLength(2);
+
+          controller.clearState();
+
+          expect(controller.state).toStrictEqual(
+            getDefaultClaimsControllerState(),
+          );
+          expect(controller.state.claims).toHaveLength(0);
+          expect(controller.state.drafts).toHaveLength(0);
+        },
+      );
+    });
+  });
 });

packages/claims-controller/src/ClaimsController.ts

@@ -292,6 +292,15 @@ export class ClaimsController extends BaseController<
     });
   }
 
+  /**
+   * Clears the claims state and resets to default values.
+   */
+  clearState(): void {
+    this.update(() => {
+      return getDefaultClaimsControllerState();
+    });
+  }
+
   /**
    * Validate the claim before submitting it.
    *

packages/shield-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Added new public method, `clearState` to clear/reset the shield controller state. ([#7780](https://github.com/MetaMask/core/pull/7780))
+
 ### Changed
 
 - Bump `@metamask/transaction-controller` from `^62.9.1` to `^62.12.0` ([#7642](https://github.com/MetaMask/core/pull/7642), [#7737](https://github.com/MetaMask/core/pull/7737), [#7760](https://github.com/MetaMask/core/pull/7760), [#7775](https://github.com/MetaMask/core/pull/7775))

packages/shield-controller/src/ShieldController.test.ts

@@ -7,9 +7,11 @@ import type { TransactionMeta } from '@metamask/transaction-controller';
 import type { TransactionControllerState } from '@metamask/transaction-controller';
 
 import {
+  getDefaultShieldControllerState,
   ShieldController,
   ShieldControllerMessenger,
 } from './ShieldController';
+import type { ShieldControllerState } from './ShieldController';
 import type { NormalizeSignatureRequestFn, ShieldBackend } from './types';
 import { TX_META_SIMULATION_DATA_MOCKS } from '../tests/data';
 import { createMockBackend, MOCK_COVERAGE_ID } from '../tests/mocks/backend';
@@ -27,16 +29,19 @@ import {
  * @param options.coverageHistoryLimit - The coverage history limit.
  * @param options.transactionHistoryLimit - The transaction history limit.
  * @param options.normalizeSignatureRequest - The function to normalize the signature request.
+ * @param options.state - The initial state for the controller.
  * @returns Objects that have been created for testing.
  */
 function setup({
   coverageHistoryLimit,
   transactionHistoryLimit,
   normalizeSignatureRequest,
+  state,
 }: {
   coverageHistoryLimit?: number;
   transactionHistoryLimit?: number;
   normalizeSignatureRequest?: NormalizeSignatureRequestFn;
+  state?: Partial<ShieldControllerState>;
 } = {}): {
   controller: ShieldController;
   messenger: ShieldControllerMessenger;
@@ -52,6 +57,7 @@ function setup({
     transactionHistoryLimit,
     messenger,
     normalizeSignatureRequest,
+    state,
   });
   controller.start();
   return {
@@ -587,4 +593,30 @@ describe('ShieldController', () => {
         `);
     });
   });
+
+  describe('clearState', () => {
+    it('should reset state to default values', () => {
+      const txMeta = generateMockTxMeta();
+      const { controller } = setup({
+        state: {
+          // @ts-expect-error - testing mock
+          coverageResults: {
+            [txMeta.id]: {
+              results: [{ status: 'covered', coverageId: MOCK_COVERAGE_ID }],
+            },
+          },
+          orderedTransactionHistory: [txMeta.id],
+        },
+      });
+
+      expect(Object.keys(controller.state.coverageResults)).toHaveLength(1);
+      expect(controller.state.orderedTransactionHistory).toHaveLength(1);
+
+      controller.clearState();
+
+      expect(controller.state).toStrictEqual(getDefaultShieldControllerState());
+      expect(Object.keys(controller.state.coverageResults)).toHaveLength(0);
+      expect(controller.state.orderedTransactionHistory).toHaveLength(0);
+    });
+  });
 });

packages/shield-controller/src/ShieldController.ts

@@ -243,6 +243,15 @@ export class ShieldController extends BaseController<
     );
   }
 
+  /**
+   * Clears the shield state and resets to default values.
+   */
+  clearState(): void {
+    this.update(() => {
+      return getDefaultShieldControllerState();
+    });
+  }
+
   #handleSignatureControllerStateChange(
     signatureRequests: Record<string, SignatureRequest>,
     previousSignatureRequests: Record<string, SignatureRequest> | undefined,

packages/subscription-controller/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added new public method `clearState` to clear/reset the subscription controller state. ([#7780](https://github.com/MetaMask/core/pull/7780))
 - Added SubscriptionController `clearLastSelectedPaymentMethod` method ([#7768](https://github.com/MetaMask/core/pull/7768))
 
 ### Changed

packages/subscription-controller/src/SubscriptionController.test.ts

@@ -1837,6 +1837,44 @@ describe('SubscriptionController', () => {
     });
   });
 
+  describe('clearState', () => {
+    it('should reset state to default values', async () => {
+      await withController(
+        {
+          state: {
+            subscriptions: [MOCK_SUBSCRIPTION],
+            pricing: MOCK_PRICE_INFO_RESPONSE,
+            lastSelectedPaymentMethod: {
+              [PRODUCT_TYPES.SHIELD]: {
+                type: PAYMENT_TYPES.byCrypto,
+                paymentTokenAddress: '0xtoken',
+                paymentTokenSymbol: 'USDT',
+                plan: RECURRING_INTERVALS.month,
+              },
+            },
+          },
+        },
+        async ({ controller }) => {
+          expect(controller.state.subscriptions).toStrictEqual([
+            MOCK_SUBSCRIPTION,
+          ]);
+          expect(controller.state.pricing).toStrictEqual(
+            MOCK_PRICE_INFO_RESPONSE,
+          );
+
+          controller.clearState();
+
+          expect(controller.state).toStrictEqual(
+            getDefaultSubscriptionControllerState(),
+          );
+          expect(controller.state.subscriptions).toHaveLength(0);
+          expect(controller.state.pricing).toBeUndefined();
+          expect(controller.state.lastSelectedPaymentMethod).toBeUndefined();
+        },
+      );
+    });
+  });
+
   describe('submitSponsorshipIntents', () => {
     const MOCK_SUBMISSION_INTENTS_REQUEST: SubmitSponsorshipIntentsMethodParams =
       {

packages/subscription-controller/src/SubscriptionController.ts

@@ -943,6 +943,16 @@ export class SubscriptionController extends StaticIntervalPollingController()<
     return tokenAmount.toFixed(0);
   }
 
+  /**
+   * Clears the subscription state and resets to default values.
+   */
+  clearState(): void {
+    const defaultState = getDefaultSubscriptionControllerState();
+    this.update(() => {
+      return defaultState;
+    });
+  }
+
   /**
    * Triggers an access token refresh.
    */