diff --git a/packages/earn-controller/package.json b/packages/earn-controller/package.json index f6aaca21a1..17c6f35dce 100644 --- a/packages/earn-controller/package.json +++ b/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", diff --git a/packages/earn-controller/src/EarnController.test.ts b/packages/earn-controller/src/EarnController.test.ts index a3e02738fd..f8232522f7 100644 --- a/packages/earn-controller/src/EarnController.test.ts +++ b/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 => { + 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); + }); + }); }); }); diff --git a/packages/earn-controller/src/EarnController.ts b/packages/earn-controller/src/EarnController.ts index e1b6d72447..970f7b919e 100644 --- a/packages/earn-controller/src/EarnController.ts +++ b/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([ + 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) { diff --git a/packages/earn-controller/tsconfig.build.json b/packages/earn-controller/tsconfig.build.json index 60df451b56..e56c9bfff2 100644 --- a/packages/earn-controller/tsconfig.build.json +++ b/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"] diff --git a/packages/earn-controller/tsconfig.json b/packages/earn-controller/tsconfig.json index bf1ccd4b0e..4b0aa0dcb8 100644 --- a/packages/earn-controller/tsconfig.json +++ b/packages/earn-controller/tsconfig.json @@ -13,6 +13,9 @@ }, { "path": "../accounts-controller" + }, + { + "path": "../transaction-controller" } ] } diff --git a/yarn.lock b/yarn.lock index ec1bcc91d5..32d1ee7e12 100644 --- a/yarn.lock +++ b/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"