Add approval controller count methods (#304)

Adds two methods to the approval controller: - `getTotalApprovalCount`, for getting the total count of all approvals, regardless of origin and type - `getApprovalCount`, for getting the count of all approvals by origin, type, or both
MetaMask · Oct 11, 2023 · 4ecc383 · 4ecc383
1 parent f71eb18
commit 4ecc383
Show file tree

Hide file tree

Showing 3 changed files with 150 additions and 7 deletions.
diff --git a/src/approval/ApprovalController.ts b/src/approval/ApprovalController.ts
@@ -140,6 +140,54 @@ export class ApprovalController extends BaseController<ApprovalConfig, ApprovalS
       : undefined;
   }
 
+  /**
+   * Gets the number of pending approvals, by origin and/or type.
+   * To only count approvals without a type, the `type` must be specified as
+   * `null`.
+   *
+   * If only `origin` is specified, all approvals for that origin will be counted,
+   * regardless of type.
+   * If only `type` is specified, all approvals for that type will be counted,
+   * regardless of origin.
+   * If both `origin` and `type` are specified, 0 or 1 will be returned.
+   *
+   * @param opts - Options bag.
+   * @param opts.origin - An approval origin.
+   * @param opts.type - The type of the approval request.
+   * @returns The pending approval count for the given origin and/or type.
+   */
+  getApprovalCount(opts: { origin?: string; type?: string | null} = {}): number {
+    if (!opts.origin && !opts.type && opts.type !== null) {
+      throw new Error('Must specify origin, type, or both.');
+    }
+    const { origin } = opts;
+    const _type = opts.type === null ? NO_TYPE : opts.type as ApprovalType;
+
+    if (origin && _type) {
+      return Number(Boolean(this._origins.get(origin)?.has(_type)));
+    }
+
+    if (origin) {
+      return this._origins.get(origin)?.size || 0;
+    }
+
+    // Only "type" was specified
+    let count = 0;
+    for (const [_origin, types] of this._origins) {
+      if (types.has(_type)) {
+        count += 1;
+      }
+    }
+    return count;
+  }
+
+  /**
+   * @returns The total pending approval count, for all types and origins.
+   */
+  getTotalApprovalCount(): number {
+    return this._approvals.size;
+  }
+
   /**
    * Checks if there's a pending approval request for the given id, or origin
    * and type pair if no id is specified.
@@ -149,18 +197,19 @@ export class ApprovalController extends BaseController<ApprovalConfig, ApprovalS
    * @param opts.id - The id of the approval request.
    * @param opts.origin - The origin of the approval request.
    * @param opts.type - The type of the approval request.
-   * @returns True if an approval is found, false otherwise.
+   * @returns `true` if an approval is found, `false` otherwise.
    */
-  has(opts: { id?: string; origin?: string; type?: string }): boolean {
+  has(opts: { id?: string; origin?: string; type?: string } = {}): boolean {
+    const { id, origin } = opts;
     const _type = opts.type === undefined ? NO_TYPE : opts.type;
     if (!_type) {
       throw new Error('May not specify falsy type.');
     }
 
-    if (opts.id) {
-      return this._approvals.has(opts.id);
-    } else if (opts.origin) {
-      return Boolean(this._origins.get(opts.origin)?.has(_type));
+    if (id) {
+      return this._approvals.has(id);
+    } else if (origin) {
+      return Boolean(this._origins.get(origin)?.has(_type));
     }
     throw new Error('Must specify id or origin.');
   }

diff --git a/tests/ApprovalController.test.js b/tests/ApprovalController.test.js
@@ -48,12 +48,23 @@ describe('ApprovalController: Input Validation', () => {
     });
   });
 
+  describe('getApprovalCount', () => {
+    it('validates input', () => {
+      const approvalController = getApprovalController();
+
+      expect(() => approvalController.getApprovalCount()).toThrow(getApprovalCountParamsError());
+      expect(() => approvalController.getApprovalCount({})).toThrow(getApprovalCountParamsError());
+      expect(() => approvalController.getApprovalCount({ origin: null })).toThrow(getApprovalCountParamsError());
+      expect(() => approvalController.getApprovalCount({ type: false })).toThrow(getApprovalCountParamsError());
+    });
+  });
+
   describe('has', () => {
     it('validates input', () => {
       const approvalController = getApprovalController();
 
+      expect(() => approvalController.has()).toThrow(getMissingIdOrOriginError());
       expect(() => approvalController.has({})).toThrow(getMissingIdOrOriginError());
-
       expect(() => approvalController.has({ type: false })).toThrow(getNoFalsyTypeError());
     });
   });
@@ -139,6 +150,10 @@ function getMissingIdOrOriginError() {
   return getError('Must specify id or origin.');
 }
 
+function getApprovalCountParamsError() {
+  return getError('Must specify origin, type, or both.');
+}
+
 function getError(message, code) {
   const err = {
     name: 'Error',

diff --git a/tests/ApprovalController.test.ts b/tests/ApprovalController.test.ts
@@ -133,6 +133,85 @@ describe('approval controller', () => {
     });
   });
 
+  describe('getApprovalCount', () => {
+    let approvalController: ApprovalController;
+    let addWithCatch: (args: any) => void;
+
+    beforeEach(() => {
+      approvalController = new ApprovalController({ ...defaultConfig });
+      addWithCatch = (args: any) => approvalController.add(args).catch(() => undefined);
+    });
+
+    it('gets the count when specifying origin and type', () => {
+      addWithCatch({ id: '1', origin: 'origin1' });
+      addWithCatch({ id: '2', origin: 'origin1', type: 'type1' });
+      addWithCatch({ id: '3', origin: 'origin2', type: 'type1' });
+
+      expect(approvalController.getApprovalCount({ origin: 'origin1', type: null })).toEqual(1);
+      expect(approvalController.getApprovalCount({ origin: 'origin1', type: 'type1' })).toEqual(1);
+      expect(approvalController.getApprovalCount({ origin: 'origin1', type: 'type2' })).toEqual(0);
+
+      expect(approvalController.getApprovalCount({ origin: 'origin2', type: null })).toEqual(0);
+      expect(approvalController.getApprovalCount({ origin: 'origin2', type: 'type1' })).toEqual(1);
+      expect(approvalController.getApprovalCount({ origin: 'origin2', type: 'type2' })).toEqual(0);
+
+      expect(approvalController.getApprovalCount({ origin: 'origin3', type: null })).toEqual(0);
+      expect(approvalController.getApprovalCount({ origin: 'origin3', type: 'type1' })).toEqual(0);
+      expect(approvalController.getApprovalCount({ origin: 'origin3', type: 'type2' })).toEqual(0);
+    });
+
+    it('gets the count when specifying origin only', () => {
+      addWithCatch({ id: '1', origin: 'origin1' });
+      addWithCatch({ id: '2', origin: 'origin1', type: 'type1' });
+      addWithCatch({ id: '3', origin: 'origin2', type: 'type1' });
+
+      expect(approvalController.getApprovalCount({ origin: 'origin1' })).toEqual(2);
+
+      expect(approvalController.getApprovalCount({ origin: 'origin2' })).toEqual(1);
+
+      expect(approvalController.getApprovalCount({ origin: 'origin3' })).toEqual(0);
+    });
+
+    it('gets the count when specifying type only', () => {
+      addWithCatch({ id: '1', origin: 'origin1' });
+      addWithCatch({ id: '2', origin: 'origin1', type: 'type1' });
+      addWithCatch({ id: '3', origin: 'origin2', type: 'type1' });
+      addWithCatch({ id: '4', origin: 'origin2', type: 'type2' });
+
+      expect(approvalController.getApprovalCount({ type: null })).toEqual(1);
+
+      expect(approvalController.getApprovalCount({ type: 'type1' })).toEqual(2);
+
+      expect(approvalController.getApprovalCount({ type: 'type2' })).toEqual(1);
+
+      expect(approvalController.getApprovalCount({ type: 'type3' })).toEqual(0);
+    });
+  });
+
+  describe('getTotalApprovalCount', () => {
+    it('gets the total approval count', () => {
+      const approvalController = new ApprovalController({ ...defaultConfig });
+      expect(approvalController.getTotalApprovalCount()).toEqual(0);
+
+      const addWithCatch = (args: any) => approvalController.add(args).catch(() => undefined);
+
+      addWithCatch({ id: '1', origin: 'origin1' });
+      expect(approvalController.getTotalApprovalCount()).toEqual(1);
+
+      addWithCatch({ id: '2', origin: 'origin1', type: 'type1' });
+      expect(approvalController.getTotalApprovalCount()).toEqual(2);
+
+      addWithCatch({ id: '3', origin: 'origin2', type: 'type1' });
+      expect(approvalController.getTotalApprovalCount()).toEqual(3);
+
+      approvalController.reject('2', new Error('foo'));
+      expect(approvalController.getTotalApprovalCount()).toEqual(2);
+
+      approvalController.clear();
+      expect(approvalController.getTotalApprovalCount()).toEqual(0);
+    });
+  });
+
   describe('has', () => {
     let approvalController: ApprovalController;