feat: implement cohort-based gating system for subscriptions

packages/subscription-controller/CHANGELOG.md

@@ -10,6 +10,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Added `lastSubscription` in state returned from `getSubscriptions` method ([#7110](https://github.com/MetaMask/core/pull/7110))
+- Add `assignUserToCohort` method to assign users to cohorts via backend API ([#7099](https://github.com/MetaMask/core/pull/7099))
+- Add cohort-related types: `Cohort`, `CohortName`, `BalanceCategory`, `AssignCohortRequest`, `GetSubscriptionsEligibilitiesRequest` ([#7099](https://github.com/MetaMask/core/pull/7099))
+- Add cohort-related constants: `COHORT_NAMES`, `BALANCE_CATEGORIES`, `SubscriptionUserEvent` ([#7099](https://github.com/MetaMask/core/pull/7099))
+- Add cohort fields to `SubscriptionEligibility` type: `cohorts`, `assignedCohort`, `hasAssignedCohortExpired` ([#7099](https://github.com/MetaMask/core/pull/7099))
+- Add `ShieldCohortAssigned` event to `SubscriptionUserEvent` ([#7099](https://github.com/MetaMask/core/pull/7099))
+- Add optional `balanceCategory` parameter to `getSubscriptionsEligibilities` for privacy-preserving balance evaluation ([#7099](https://github.com/MetaMask/core/pull/7099))
+- Add optional `cohort` field to `SubmitUserEventRequest` for event tracking ([#7099](https://github.com/MetaMask/core/pull/7099))
+
+### Changed
+
+- Refactor `SubscriptionService.makeRequest` to accept query parameters for cleaner URL construction ([#7099](https://github.com/MetaMask/core/pull/7099))
 
 ## [3.3.0]
 

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

@@ -135,6 +135,21 @@ const MOCK_GET_SUBSCRIPTIONS_RESPONSE = {
   trialedProducts: [],
 };
 
+const MOCK_COHORTS = [
+  {
+    cohort: 'post_tx',
+    eligibilityRate: 0.8,
+    priority: 1,
+    eligible: true,
+  },
+  {
+    cohort: 'wallet_home',
+    eligibilityRate: 0.2,
+    priority: 2,
+    eligible: true,
+  },
+];
+
 /**
  * Creates a custom subscription messenger, in case tests need different permissions
  *
@@ -219,6 +234,7 @@ function createMockSubscriptionService() {
   const mockGetSubscriptionsEligibilities = jest.fn();
   const mockSubmitUserEvent = jest.fn();
   const mockSubmitSponsorshipIntents = jest.fn();
+  const mockAssignUserToCohort = jest.fn();
 
   const mockService = {
     getSubscriptions: mockGetSubscriptions,
@@ -233,6 +249,7 @@ function createMockSubscriptionService() {
     getSubscriptionsEligibilities: mockGetSubscriptionsEligibilities,
     submitUserEvent: mockSubmitUserEvent,
     submitSponsorshipIntents: mockSubmitSponsorshipIntents,
+    assignUserToCohort: mockAssignUserToCohort,
   };
 
   return {
@@ -246,6 +263,7 @@ function createMockSubscriptionService() {
     mockUpdatePaymentMethodCard,
     mockUpdatePaymentMethodCrypto,
     mockSubmitSponsorshipIntents,
+    mockAssignUserToCohort,
   };
 }
 
@@ -1408,6 +1426,9 @@ describe('SubscriptionController', () => {
       canSubscribe: true,
       minBalanceUSD: 100,
       canViewEntryModal: true,
+      cohorts: [],
+      assignedCohort: null,
+      hasAssignedCohortExpired: false,
     };
 
     it('should get the subscriptions eligibilities', async () => {
@@ -1421,6 +1442,29 @@ describe('SubscriptionController', () => {
       });
     });
 
+    it('should get the subscriptions eligibilities with balanceCategory parameter', async () => {
+      await withController(async ({ controller, mockService }) => {
+        const mockEligibilityWithCohorts: SubscriptionEligibility = {
+          ...MOCK_SUBSCRIPTION_ELIGIBILITY,
+          cohorts: MOCK_COHORTS,
+          assignedCohort: 'post_tx',
+        };
+
+        mockService.getSubscriptionsEligibilities.mockResolvedValue([
+          mockEligibilityWithCohorts,
+        ]);
+
+        const balanceCategory = '1k-9.9k';
+        const result = await controller.getSubscriptionsEligibilities({
+          balanceCategory,
+        });
+        expect(result).toStrictEqual([mockEligibilityWithCohorts]);
+        expect(mockService.getSubscriptionsEligibilities).toHaveBeenCalledWith({
+          balanceCategory,
+        });
+      });
+    });
+
     it('should handle subscription service errors', async () => {
       await withController(async ({ controller, mockService }) => {
         const errorMessage = 'Failed to get subscriptions eligibilities';
@@ -1449,6 +1493,26 @@ describe('SubscriptionController', () => {
         expect(submitUserEventSpy).toHaveBeenCalledWith({
           event: SubscriptionUserEvent.ShieldEntryModalViewed,
         });
+        expect(submitUserEventSpy).toHaveBeenCalledTimes(1);
+      });
+    });
+
+    it('should submit user event with cohort successfully', async () => {
+      await withController(async ({ controller, mockService }) => {
+        const submitUserEventSpy = jest
+          .spyOn(mockService, 'submitUserEvent')
+          .mockResolvedValue(undefined);
+
+        const result = await controller.submitUserEvent({
+          event: SubscriptionUserEvent.ShieldCohortAssigned,
+          cohort: 'post_tx',
+        });
+        expect(result).toBeUndefined();
+        expect(submitUserEventSpy).toHaveBeenCalledWith({
+          event: SubscriptionUserEvent.ShieldCohortAssigned,
+          cohort: 'post_tx',
+        });
+        expect(submitUserEventSpy).toHaveBeenCalledTimes(1);
       });
     });
 
@@ -1468,6 +1532,38 @@ describe('SubscriptionController', () => {
     });
   });
 
+  describe('assignUserToCohort', () => {
+    it('should assign user to cohort successfully', async () => {
+      await withController(async ({ controller, mockService }) => {
+        const assignUserToCohortSpy = jest
+          .spyOn(mockService, 'assignUserToCohort')
+          .mockResolvedValue(undefined);
+
+        const result = await controller.assignUserToCohort({
+          cohort: 'post_tx',
+        });
+        expect(result).toBeUndefined();
+        expect(assignUserToCohortSpy).toHaveBeenCalledWith({
+          cohort: 'post_tx',
+        });
+        expect(assignUserToCohortSpy).toHaveBeenCalledTimes(1);
+      });
+    });
+
+    it('should handle subscription service errors', async () => {
+      await withController(async ({ controller, mockService }) => {
+        const errorMessage = 'Failed to assign user to cohort';
+        mockService.assignUserToCohort.mockRejectedValue(
+          new SubscriptionServiceError(errorMessage),
+        );
+
+        await expect(
+          controller.assignUserToCohort({ cohort: 'post_tx' }),
+        ).rejects.toThrow(SubscriptionServiceError);
+      });
+    });
+  });
+
   describe('cacheLastSelectedPaymentMethod', () => {
     const MOCK_CACHED_PAYMENT_METHOD: CachedLastSelectedPaymentMethod = {
       type: PAYMENT_TYPES.byCrypto,

packages/subscription-controller/src/SubscriptionController.ts

@@ -18,9 +18,11 @@ import {
   SubscriptionControllerErrorMessage,
 } from './constants';
 import type {
+  AssignCohortRequest,
   BillingPortalResponse,
   GetCryptoApproveTransactionRequest,
   GetCryptoApproveTransactionResponse,
+  GetSubscriptionsEligibilitiesRequest,
   ProductPrice,
   SubscriptionEligibility,
   StartCryptoSubscriptionRequest,
@@ -410,10 +412,15 @@ export class SubscriptionController extends StaticIntervalPollingController()<
   /**
    * Get the subscriptions eligibilities.
    *
+   * @param request - Optional request object containing user balance to check cohort eligibility.
    * @returns The subscriptions eligibilities.
    */
-  async getSubscriptionsEligibilities(): Promise<SubscriptionEligibility[]> {
-    return await this.#subscriptionService.getSubscriptionsEligibilities();
+  async getSubscriptionsEligibilities(
+    request?: GetSubscriptionsEligibilitiesRequest,
+  ): Promise<SubscriptionEligibility[]> {
+    return await this.#subscriptionService.getSubscriptionsEligibilities(
+      request,
+    );
   }
 
   async cancelSubscription(request: { subscriptionId: string }) {
@@ -704,12 +711,22 @@ export class SubscriptionController extends StaticIntervalPollingController()<
    * Submit a user event from the UI. (e.g. shield modal viewed)
    *
    * @param request - Request object containing the event to submit.
-   * @example { event: SubscriptionUserEvent.ShieldEntryModalViewed }
+   * @example { event: SubscriptionUserEvent.ShieldEntryModalViewed, cohort: 'post_tx' }
    */
   async submitUserEvent(request: SubmitUserEventRequest) {
     await this.#subscriptionService.submitUserEvent(request);
   }
 
+  /**
+   * Assign user to a cohort.
+   *
+   * @param request - Request object containing the cohort to assign the user to.
+   * @example { cohort: 'post_tx' }
+   */
+  async assignUserToCohort(request: AssignCohortRequest): Promise<void> {
+    await this.#subscriptionService.assignUserToCohort(request);
+  }
+
   async _executePoll(): Promise<void> {
     await this.getSubscriptions();
     if (this.#shouldCallRefreshAuthToken) {