feat: STAKE-1005: refresh staking data when staking txs are confirmed

packages/earn-controller/package.json

@@ -56,6 +56,7 @@
     "@metamask/accounts-controller": "^27.0.0",
     "@metamask/auto-changelog": "^3.4.4",
     "@metamask/network-controller": "^23.2.0",
+    "@metamask/transaction-controller": "^53.0.0",
     "@types/jest": "^27.4.1",
     "deepmerge": "^4.2.2",
     "jest": "^27.5.1",

packages/earn-controller/src/EarnController.test.ts

@@ -1,8 +1,13 @@
 import type { AccountsController } from '@metamask/accounts-controller';
 import { Messenger } from '@metamask/base-controller';
+import { toHex } from '@metamask/controller-utils';
 import { getDefaultNetworkControllerState } from '@metamask/network-controller';
 import { StakeSdk, StakingApiService } from '@metamask/stake-sdk';
 
+import type {
+  EarnControllerGetStateAction,
+  EarnControllerStateChangeEvent,
+} from './EarnController';
 import {
   EarnController,
   type EarnControllerState,
@@ -13,6 +18,11 @@ import {
   type AllowedActions,
   type AllowedEvents,
 } from './EarnController';
+import type { TransactionMeta } from '../../transaction-controller/src';
+import {
+  TransactionStatus,
+  TransactionType,
+} from '../../transaction-controller/src';
 
 jest.mock('@metamask/stake-sdk', () => ({
   StakeSdk: {
@@ -63,6 +73,7 @@ function getEarnControllerMessenger(
     allowedEvents: [
       'NetworkController:stateChange',
       'AccountsController:selectedAccountChange',
+      'TransactionController:transactionConfirmed',
     ],
   });
 }
@@ -102,6 +113,30 @@ const mockAccount1Address = '0x1234';
 
 const mockAccount2Address = '0xabc';
 
+const createMockTransaction = ({
+  id = '1',
+  type = TransactionType.stakingDeposit,
+  chainId = toHex(1),
+  networkClientId = 'networkClientIdMock',
+  time = 123456789,
+  status = TransactionStatus.confirmed,
+  txParams = {
+    gasUsed: '0x5208',
+    from: mockAccount1Address,
+    to: mockAccount2Address,
+  },
+}: Partial<TransactionMeta> = {}): TransactionMeta => {
+  return {
+    id,
+    type,
+    chainId,
+    networkClientId,
+    time,
+    status,
+    txParams,
+  };
+};
+
 const mockPooledStakes = {
   account: mockAccount1Address,
   lifetimeRewards: '100',
@@ -634,5 +669,86 @@ describe('EarnController', () => {
         });
       });
     });
+
+    describe('On transaction confirmed', () => {
+      let controller: EarnController;
+      let messenger: Messenger<
+        EarnControllerGetStateAction | AllowedActions,
+        EarnControllerStateChangeEvent | AllowedEvents
+      >;
+
+      beforeEach(() => {
+        const earnController = setupController();
+        controller = earnController.controller;
+        messenger = earnController.messenger;
+
+        jest.spyOn(controller, 'refreshPooledStakes').mockResolvedValue();
+      });
+
+      it('updates pooled stakes for staking deposit transaction type', () => {
+        const MOCK_CONFIRMED_DEPOSIT_TX = createMockTransaction({
+          type: TransactionType.stakingDeposit,
+          status: TransactionStatus.confirmed,
+        });
+
+        messenger.publish(
+          'TransactionController:transactionConfirmed',
+          MOCK_CONFIRMED_DEPOSIT_TX,
+        );
+
+        expect(controller.refreshPooledStakes).toHaveBeenNthCalledWith(1, {
+          address: MOCK_CONFIRMED_DEPOSIT_TX.txParams.from,
+          resetCache: true,
+        });
+      });
+
+      it('updates pooled stakes for staking unstake transaction type', () => {
+        const MOCK_CONFIRMED_UNSTAKE_TX = createMockTransaction({
+          type: TransactionType.stakingUnstake,
+          status: TransactionStatus.confirmed,
+        });
+
+        messenger.publish(
+          'TransactionController:transactionConfirmed',
+          MOCK_CONFIRMED_UNSTAKE_TX,
+        );
+
+        expect(controller.refreshPooledStakes).toHaveBeenNthCalledWith(1, {
+          address: MOCK_CONFIRMED_UNSTAKE_TX.txParams.from,
+          resetCache: true,
+        });
+      });
+
+      it('updates pooled stakes for staking claim transaction type', () => {
+        const MOCK_CONFIRMED_CLAIM_TX = createMockTransaction({
+          type: TransactionType.stakingClaim,
+          status: TransactionStatus.confirmed,
+        });
+
+        messenger.publish(
+          'TransactionController:transactionConfirmed',
+          MOCK_CONFIRMED_CLAIM_TX,
+        );
+
+        expect(controller.refreshPooledStakes).toHaveBeenNthCalledWith(1, {
+          address: MOCK_CONFIRMED_CLAIM_TX.txParams.from,
+          resetCache: true,
+        });
+      });
+
+      it('ignores non-staking transaction types', () => {
+        const MOCK_CONFIRMED_SWAP_TX = createMockTransaction({
+          type: TransactionType.swap,
+          status: TransactionStatus.confirmed,
+        });
+
+        messenger.publish(
+          'TransactionController:transactionConfirmed',
+          MOCK_CONFIRMED_SWAP_TX,
+        );
+
+        expect(controller.refreshPooledStakes).toHaveBeenCalledTimes(0);
+      });
+    });
   });
 });

packages/earn-controller/src/EarnController.ts

@@ -25,6 +25,10 @@ import {
   type VaultDailyApy,
   type VaultApyAverages,
 } from '@metamask/stake-sdk';
+import {
+  TransactionType,
+  type TransactionControllerTransactionConfirmedEvent,
+} from '@metamask/transaction-controller';
 
 import type {
   RefreshPooledStakesOptions,
@@ -58,6 +62,17 @@ export type StablecoinVault = {
   liquidity: string;
 };
 
+type StakingTransactionTypes =
+  | TransactionType.stakingDeposit
+  | TransactionType.stakingUnstake
+  | TransactionType.stakingClaim;
+
+const stakingTransactionTypes = new Set<StakingTransactionTypes>([
+  TransactionType.stakingDeposit,
+  TransactionType.stakingUnstake,
+  TransactionType.stakingClaim,
+]);
+
 /**
  * Metadata for the EarnController.
  */
@@ -178,7 +193,8 @@ export type EarnControllerEvents = EarnControllerStateChangeEvent;
  */
 export type AllowedEvents =
   | AccountsControllerSelectedAccountChangeEvent
-  | NetworkControllerStateChangeEvent;
+  | NetworkControllerStateChangeEvent
+  | TransactionControllerTransactionConfirmedEvent;
 
 /**
  * The messenger which is restricted to actions and events accessed by
@@ -233,6 +249,7 @@ export class EarnController extends BaseController<
     );
     this.#selectedNetworkClientId = selectedNetworkClientId;
 
+    // Listen for network changes
     this.messagingSystem.subscribe(
       'NetworkController:stateChange',
       (networkControllerState) => {
@@ -265,6 +282,30 @@ export class EarnController extends BaseController<
         this.refreshPooledStakes({ address }).catch(console.error);
       },
     );
+
+    // Listen for confirmed staking transactions
+    this.messagingSystem.subscribe(
+      'TransactionController:transactionConfirmed',
+      (transactionMeta) => {
+        /**
+         * When we speed up a transaction, we set the type as Retry and we lose
+         * information about type of transaction that is being set up, so we use
+         * original type to track that information.
+         */
+        const { type, originalType } = transactionMeta;
+
+        const isStakingTransaction =
+          stakingTransactionTypes.has(type as StakingTransactionTypes) ||
+          stakingTransactionTypes.has(originalType as StakingTransactionTypes);
+
+        if (isStakingTransaction) {
+          const sender = transactionMeta.txParams.from;
+          this.refreshPooledStakes({ resetCache: true, address: sender }).catch(
+            console.error,
+          );
+        }
+      },
+    );
   }
 
   #initializeSDK(networkClientId?: string) {

packages/earn-controller/tsconfig.build.json

@@ -14,6 +14,9 @@
     },
     {
       "path": "../accounts-controller/tsconfig.build.json"
+    },
+    {
+      "path": "../transaction-controller/tsconfig.build.json"
     }
   ],
   "include": ["../../types", "./src"]

packages/earn-controller/tsconfig.json

@@ -13,6 +13,9 @@
     },
     {
       "path": "../accounts-controller"
+    },
+    {
+      "path": "../transaction-controller"
     }
   ]
 }

yarn.lock

@@ -2962,6 +2962,7 @@ __metadata:
     "@metamask/controller-utils": "npm:^11.7.0"
     "@metamask/network-controller": "npm:^23.2.0"
     "@metamask/stake-sdk": "npm:^1.0.0"
+    "@metamask/transaction-controller": "npm:^53.0.0"
     "@types/jest": "npm:^27.4.1"
     deepmerge: "npm:^4.2.2"
     jest: "npm:^27.5.1"