feat: 🎸 add subscription to staking getters

allow users to subscribe to an account's controller, nominations and
eraInfo

src/api/client/Staking.ts

@@ -1,3 +1,5 @@
+import { Option, u32, u128 } from '@polkadot/types';
+import { PalletStakingActiveEraInfo } from '@polkadot/types/lookup';
 import BigNumber from 'bignumber.js';
 
 import { Account, Context } from '~/internal';
@@ -7,6 +9,8 @@ import {
   ResultSet,
   StakingCommission,
   StakingEraInfo,
+  SubCallback,
+  UnsubCallback,
 } from '~/types';
 import {
   accountIdToString,
@@ -81,16 +85,113 @@ export class Staking {
   /**
    * Retrieve the current staking era
    *
-   * TODO support subscription?
-   * TODO bundle more info?
+   * @note can be subscribed to, if connected to node using a web socket
    */
-  public async eraInfo(): Promise<StakingEraInfo> {
+  public async eraInfo(): Promise<StakingEraInfo>;
+  public async eraInfo(callback: SubCallback<StakingEraInfo>): Promise<UnsubCallback>;
+
+  // eslint-disable-next-line require-jsdoc
+  public async eraInfo(
+    callback?: SubCallback<StakingEraInfo>
+  ): Promise<StakingEraInfo | UnsubCallback> {
     const {
       context: {
         polymeshApi: { query },
       },
+      context,
     } = this;
 
+    const assembleResult = (
+      rawActiveEra: Option<PalletStakingActiveEraInfo>,
+      rawCurrentEra: Option<u32>,
+      rawPlannedSession: u32,
+      rawTotalStaked: u128
+    ): StakingEraInfo => {
+      let activeEra: ActiveEraInfo;
+      if (rawActiveEra.isNone) {
+        activeEra = { index: new BigNumber(0), start: new BigNumber(0) };
+      } else {
+        activeEra = activeEraStakingToActiveEraInfo(rawActiveEra.unwrap());
+      }
+
+      let currentEra: BigNumber;
+      if (rawCurrentEra.isNone) {
+        currentEra = new BigNumber(0);
+      } else {
+        currentEra = u32ToBigNumber(rawCurrentEra.unwrap());
+      }
+
+      const plannedSession = u32ToBigNumber(rawPlannedSession);
+      const totalStaked = u128ToBigNumber(rawTotalStaked);
+
+      return {
+        activeEra: activeEra.index,
+        activeEraStart: activeEra.start,
+        currentEra,
+        plannedSession,
+        totalStaked,
+      };
+    };
+
+    if (callback) {
+      context.assertSupportsSubscription();
+
+      let rawActiveEra: Option<PalletStakingActiveEraInfo>;
+      let rawCurrentEra: Option<u32> = context.createType('Option<u32>', undefined); // workaround "no use before defined" rule
+      let rawPlannedSession: u32;
+      let rawTotalStaked: u128;
+
+      let initialized = false;
+
+      const callCb = (): void => {
+        if (!initialized) {
+          return;
+        }
+
+        const result = assembleResult(
+          rawActiveEra,
+          rawCurrentEra,
+          rawPlannedSession,
+          rawTotalStaked
+        );
+
+        // eslint-disable-next-line @typescript-eslint/no-floating-promises -- callback errors should be handled by the caller
+        callback(result);
+      };
+
+      const [activeUnsub, currentUnsub, plannedUnsub] = await Promise.all([
+        query.staking.activeEra(activeEra => {
+          rawActiveEra = activeEra;
+
+          callCb();
+        }),
+        query.staking.currentEra(async currentEra => {
+          rawCurrentEra = currentEra;
+          rawTotalStaked = await query.staking.erasTotalStake(rawCurrentEra.unwrapOr(0));
+
+          callCb();
+        }),
+        query.staking.currentPlannedSession(plannedSession => {
+          rawPlannedSession = plannedSession;
+
+          callCb();
+        }),
+      ]);
+
+      rawTotalStaked = await query.staking.erasTotalStake(rawCurrentEra.unwrapOr(0));
+
+      const unsub = (): void => {
+        activeUnsub();
+        currentUnsub();
+        plannedUnsub();
+      };
+
+      initialized = true;
+      callCb();
+
+      return unsub;
+    }
+
     const [rawActiveEra, rawCurrentEra, rawPlannedSession] = await Promise.all([
       query.staking.activeEra(),
       query.staking.currentEra(),
@@ -99,29 +200,6 @@ export class Staking {
 
     const rawTotalStaked = await query.staking.erasTotalStake(rawCurrentEra.unwrapOr(0));
 
-    let activeEra: ActiveEraInfo;
-    if (rawActiveEra.isNone) {
-      activeEra = { index: new BigNumber(0), start: new BigNumber(0) };
-    } else {
-      activeEra = activeEraStakingToActiveEraInfo(rawActiveEra.unwrap());
-    }
-
-    let currentEra: BigNumber;
-    if (rawCurrentEra.isNone) {
-      currentEra = new BigNumber(0);
-    } else {
-      currentEra = u32ToBigNumber(rawCurrentEra.unwrap());
-    }
-
-    const plannedSession = u32ToBigNumber(rawPlannedSession);
-    const totalStaked = u128ToBigNumber(rawTotalStaked);
-
-    return {
-      activeEra: activeEra.index,
-      activeEraStart: activeEra.start,
-      currentEra,
-      plannedSession,
-      totalStaked,
-    };
+    return assembleResult(rawActiveEra, rawCurrentEra, rawPlannedSession, rawTotalStaked);
   }
 }

src/api/client/__tests__/Staking.ts

@@ -1,4 +1,6 @@
+import { Option, u32, u128 } from '@polkadot/types';
 import { AccountId } from '@polkadot/types/interfaces';
+import { PalletStakingActiveEraInfo } from '@polkadot/types/lookup';
 import BigNumber from 'bignumber.js';
 import { when } from 'jest-when';
 
@@ -89,27 +91,42 @@ describe('Staking Class', () => {
 
     let rawActiveStart;
     let rawActiveIndex;
+    let rawActiveEra: Option<PalletStakingActiveEraInfo>;
+    let rawCurrentEra: Option<u32>;
+    let rawPlannedSession: u32;
+    let rawTotal: u128;
+
+    let activeEraQueryMock: jest.SpyInstance;
+    let currentEraQueryMock: jest.SpyInstance;
+    let plannedSessionQueryMock: jest.SpyInstance;
 
     beforeEach(() => {
       rawActiveIndex = dsMockUtils.createMockU32(activeIndex);
       rawActiveStart = dsMockUtils.createMockOption(dsMockUtils.createMockU64(activeStart));
-      dsMockUtils.createMockAccountData();
-      dsMockUtils.createQueryMock('staking', 'activeEra', {
-        returnValue: dsMockUtils.createMockOption(
-          dsMockUtils.createMockActiveEraInfo({ index: rawActiveIndex, start: rawActiveStart })
-        ),
+      rawActiveEra = dsMockUtils.createMockOption(
+        dsMockUtils.createMockActiveEraInfo({
+          index: rawActiveIndex,
+          start: rawActiveStart,
+        })
+      );
+      rawCurrentEra = dsMockUtils.createMockOption(dsMockUtils.createMockU32(new BigNumber(2)));
+      rawPlannedSession = dsMockUtils.createMockU32(new BigNumber(3));
+      rawTotal = dsMockUtils.createMockU128(new BigNumber(1000));
+
+      activeEraQueryMock = dsMockUtils.createQueryMock('staking', 'activeEra', {
+        returnValue: rawActiveEra,
       });
 
-      dsMockUtils.createQueryMock('staking', 'currentEra', {
-        returnValue: dsMockUtils.createMockOption(dsMockUtils.createMockU32(new BigNumber(2))),
+      currentEraQueryMock = dsMockUtils.createQueryMock('staking', 'currentEra', {
+        returnValue: rawCurrentEra,
       });
 
-      dsMockUtils.createQueryMock('staking', 'currentPlannedSession', {
-        returnValue: dsMockUtils.createMockU32(new BigNumber(3)),
+      plannedSessionQueryMock = dsMockUtils.createQueryMock('staking', 'currentPlannedSession', {
+        returnValue: rawPlannedSession,
       });
 
       dsMockUtils.createQueryMock('staking', 'erasTotalStake', {
-        returnValue: dsMockUtils.createMockU128(new BigNumber(1000)),
+        returnValue: rawTotal,
       });
     });
 
@@ -144,5 +161,50 @@ describe('Staking Class', () => {
         totalStaked: new BigNumber(1000),
       });
     });
+
+    it('should handle subscription', async () => {
+      const activeUnsub = jest.fn();
+      const eraUnsub = jest.fn();
+      const sessionUnsub = jest.fn();
+
+      type CallbackSig = (arg: unknown) => void;
+
+      activeEraQueryMock.mockImplementation((cb: CallbackSig) => {
+        cb(rawActiveEra);
+
+        return activeUnsub;
+      });
+
+      currentEraQueryMock.mockImplementation((cb: CallbackSig) => {
+        cb(rawCurrentEra);
+
+        return eraUnsub;
+      });
+
+      plannedSessionQueryMock.mockImplementation((cb: CallbackSig) => {
+        cb(rawPlannedSession);
+
+        return sessionUnsub;
+      });
+
+      const callback = jest.fn();
+
+      const result = await staking.eraInfo(callback);
+
+      expect(callback).toHaveBeenCalled();
+      expect(result).toBeInstanceOf(Function);
+
+      // ensure all unsub functions have not been called
+      expect(activeUnsub).not.toHaveBeenCalled();
+      expect(eraUnsub).not.toHaveBeenCalled();
+      expect(sessionUnsub).not.toHaveBeenCalled();
+
+      expect(() => result()).not.toThrow();
+
+      // ensure all unsub functions have been called now that unsub ran
+      expect(activeUnsub).toHaveBeenCalled();
+      expect(eraUnsub).toHaveBeenCalled();
+      expect(sessionUnsub).toHaveBeenCalled();
+    });
   });
 });