From b8218b628bb98a97bfaf2d44bac13af2707ef9eb Mon Sep 17 00:00:00 2001
From: Matthew Grainger <46547583+Matt561@users.noreply.github.com>
Date: Fri, 14 Feb 2025 14:35:05 -0500
Subject: [PATCH] feat: EarnController add resetCache arg to
 stakingApiService.getPooledStakes() (#5334)

## Explanation
This PR adds the resetCache` arg when calling
`stakingApiService.getPooledStakes` inside the `EarnController`.
<!--
Thanks for your contribution! Take a moment to answer these questions so
that reviewers have the information they need to properly understand
your changes:

* What is the current state of things and why does it need to change?
* What is the solution your changes offer and how does it work?
* Are there any changes whose purpose might not obvious to those
unfamiliar with the domain?
* If your primary goal was to update one package but you found you had
to update another one along the way, why did you do so?
* If you had to upgrade a dependency, why did you do so?
-->

## Changelog

<!--
If you're making any consumer-facing changes, list those changes here as
if you were updating a changelog, using the template below as a guide.

(CATEGORY is one of BREAKING, ADDED, CHANGED, DEPRECATED, REMOVED, or
FIXED. For security-related issues, follow the Security Advisory
process.)

Please take care to name the exact pieces of the API you've added or
changed (e.g. types, interfaces, functions, or methods).

If there are any breaking changes, make sure to offer a solution for
consumers to follow once they upgrade to the changes.

Finally, if you're only making changes to development scripts or tests,
you may replace the template below with "None".
-->

### `@metamask/earn-controller`

- **CHANGED**: Updated `refreshPooledStakes` internal call to
`stakingApiService.getPooledStakes` to force cache reset

## Checklist

- [x] I've updated the test suite for new or updated code as appropriate
- [ ] I've updated documentation (JSDoc, Markdown, etc.) for new or
updated code as appropriate
- [ ] I've highlighted breaking changes using the "BREAKING" category
above as appropriate
- [ ] I've prepared draft pull requests for clients and consumer
packages to resolve any breaking changes
---
 .../src/EarnController.test.ts                | 69 +++++++++++++++++++
 .../earn-controller/src/EarnController.ts     | 11 +--
 2 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/packages/earn-controller/src/EarnController.test.ts b/packages/earn-controller/src/EarnController.test.ts
index 3b0ab5ca796..a554be13213 100644
--- a/packages/earn-controller/src/EarnController.test.ts
+++ b/packages/earn-controller/src/EarnController.test.ts
@@ -168,6 +168,8 @@ let mockedStakingApiService: Partial<StakingApiService>;
 
 describe('EarnController', () => {
   beforeEach(() => {
+    jest.clearAllMocks();
+
     // Apply StakeSdk mock before initializing EarnController
     (StakeSdk.create as jest.Mock).mockImplementation(() => ({
       pooledStakingContract: {
@@ -292,6 +294,32 @@ describe('EarnController', () => {
       expect(controller.state.lastUpdated).toBeDefined();
     });
 
+    it('does not invalidate cache when refreshing state', async () => {
+      const { controller } = setupController();
+      await controller.refreshPooledStakingData();
+
+      expect(mockedStakingApiService.getPooledStakes).toHaveBeenNthCalledWith(
+        // First call occurs during setupController()
+        2,
+        ['0x1234'],
+        1,
+        false,
+      );
+    });
+
+    it('invalidates cache when refreshing state', async () => {
+      const { controller } = setupController();
+      await controller.refreshPooledStakingData(true);
+
+      expect(mockedStakingApiService.getPooledStakes).toHaveBeenNthCalledWith(
+        // First call occurs during setupController()
+        2,
+        ['0x1234'],
+        1,
+        true,
+      );
+    });
+
     it('handles API errors gracefully', async () => {
       const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation();
       mockedStakingApiService = {
@@ -338,6 +366,47 @@ describe('EarnController', () => {
     });
   });
 
+  describe('refreshPooledStakes', () => {
+    it('fetches without resetting cache when resetCache is false', async () => {
+      const { controller } = setupController();
+      await controller.refreshPooledStakes(false);
+
+      // Assertion on second call since the first is part of controller setup.
+      expect(mockedStakingApiService.getPooledStakes).toHaveBeenNthCalledWith(
+        2,
+        ['0x1234'],
+        1,
+        false,
+      );
+    });
+
+    it('fetches without resetting cache when resetCache is undefined', async () => {
+      const { controller } = setupController();
+      await controller.refreshPooledStakes();
+
+      // Assertion on second call since the first is part of controller setup.
+      expect(mockedStakingApiService.getPooledStakes).toHaveBeenNthCalledWith(
+        2,
+        ['0x1234'],
+        1,
+        false,
+      );
+    });
+
+    it('fetches while resetting cache', async () => {
+      const { controller } = setupController();
+      await controller.refreshPooledStakes(true);
+
+      // Assertion on second call since the first is part of controller setup.
+      expect(mockedStakingApiService.getPooledStakes).toHaveBeenNthCalledWith(
+        2,
+        ['0x1234'],
+        1,
+        true,
+      );
+    });
+  });
+
   describe('subscription handlers', () => {
     const firstAccount = createMockInternalAccount({
       address: '0x1234',
diff --git a/packages/earn-controller/src/EarnController.ts b/packages/earn-controller/src/EarnController.ts
index 5ce6e71280e..e97ac1c937d 100644
--- a/packages/earn-controller/src/EarnController.ts
+++ b/packages/earn-controller/src/EarnController.ts
@@ -11,10 +11,10 @@ import type {
 } from '@metamask/base-controller';
 import { BaseController } from '@metamask/base-controller';
 import { convertHexToDecimal } from '@metamask/controller-utils';
-import type { NetworkControllerStateChangeEvent } from '@metamask/network-controller';
 import type {
   NetworkControllerGetNetworkClientByIdAction,
   NetworkControllerGetStateAction,
+  NetworkControllerStateChangeEvent,
 } from '@metamask/network-controller';
 import {
   StakeSdk,
@@ -298,9 +298,10 @@ export class EarnController extends BaseController<
    * Fetches updated stake information including lifetime rewards, assets, and exit requests
    * from the staking API service and updates the state.
    *
+   * @param resetCache - Control whether the BE cache should be invalidated.
    * @returns A promise that resolves when the stakes data has been updated
    */
-  async refreshPooledStakes(): Promise<void> {
+  async refreshPooledStakes(resetCache = false): Promise<void> {
     const currentAccount = this.#getCurrentAccount();
     if (!currentAccount?.address) {
       return;
@@ -312,6 +313,7 @@ export class EarnController extends BaseController<
       await this.#stakingApiService.getPooledStakes(
         [currentAccount.address],
         chainId,
+        resetCache,
       );
 
     this.update((state) => {
@@ -363,14 +365,15 @@ export class EarnController extends BaseController<
    * This method allows partial success, meaning some data may update while other requests fail.
    * All errors are collected and thrown as a single error message.
    *
+   * @param resetCache - Control whether the BE cache should be invalidated.
    * @returns A promise that resolves when all possible data has been updated
    * @throws {Error} If any of the refresh operations fail, with concatenated error messages
    */
-  async refreshPooledStakingData(): Promise<void> {
+  async refreshPooledStakingData(resetCache = false): Promise<void> {
     const errors: Error[] = [];
 
     await Promise.all([
-      this.refreshPooledStakes().catch((error) => {
+      this.refreshPooledStakes(resetCache).catch((error) => {
         errors.push(error);
       }),
       this.refreshStakingEligibility().catch((error) => {