test: add test to ensure cache is loaded only once during multiple fetchTokenList calls

Author: sahar-fehri

packages/assets-controllers/src/TokenListController.test.ts

@@ -2069,6 +2069,107 @@ describe('TokenListController', () => {
       consoleErrorSpy.mockRestore();
       controller.destroy();
     });
+
+    it('should only load cache from storage once even when fetchTokenList is called multiple times', async () => {
+      // Pre-populate storage with cached data
+      const chainData: DataCache = {
+        data: sampleMainnetTokensChainsCache,
+        timestamp: Date.now(),
+      };
+      mockStorage.set(
+        `TokenListController:tokensChainsCache:${ChainId.mainnet}`,
+        chainData,
+      );
+
+      // Track how many times getItem is called
+      let getItemCallCount = 0;
+      let getAllKeysCallCount = 0;
+
+      const trackingMessenger = new Messenger({
+        namespace: MOCK_ANY_NAMESPACE,
+      });
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (trackingMessenger as any).registerActionHandler(
+        'StorageService:getItem',
+        (controllerNamespace: string, key: string) => {
+          getItemCallCount += 1;
+          const storageKey = `${controllerNamespace}:${key}`;
+          const value = mockStorage.get(storageKey);
+          return value ? { result: value } : {};
+        },
+      );
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (trackingMessenger as any).registerActionHandler(
+        'StorageService:setItem',
+        (controllerNamespace: string, key: string, value: unknown) => {
+          const storageKey = `${controllerNamespace}:${key}`;
+          mockStorage.set(storageKey, value);
+        },
+      );
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (trackingMessenger as any).registerActionHandler(
+        'StorageService:removeItem',
+        (controllerNamespace: string, key: string) => {
+          const storageKey = `${controllerNamespace}:${key}`;
+          mockStorage.delete(storageKey);
+        },
+      );
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (trackingMessenger as any).registerActionHandler(
+        'StorageService:getAllKeys',
+        (controllerNamespace: string) => {
+          getAllKeysCallCount += 1;
+          const keys: string[] = [];
+          const prefix = `${controllerNamespace}:`;
+          mockStorage.forEach((_value, key) => {
+            if (key.startsWith(prefix)) {
+              const keyWithoutNamespace = key.substring(prefix.length);
+              keys.push(keyWithoutNamespace);
+            }
+          });
+          return keys;
+        },
+      );
+
+      const restrictedMessenger = getRestrictedMessenger(trackingMessenger);
+
+      const controller = new TokenListController({
+        chainId: ChainId.mainnet,
+        messenger: restrictedMessenger,
+      });
+
+      // Wait for initialization to complete
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      // Record call counts after initialization
+      const getItemCallsAfterInit = getItemCallCount;
+      const getAllKeysCallsAfterInit = getAllKeysCallCount;
+
+      // getAllKeys should be called twice during init (once for load, once for migration check)
+      expect(getAllKeysCallsAfterInit).toBe(2);
+      // getItem should be called once for the cached chain during load
+      expect(getItemCallsAfterInit).toBe(1);
+
+      // Now call fetchTokenList multiple times
+      nock(tokenService.TOKEN_END_POINT_API)
+        .get(getTokensPath(ChainId.mainnet))
+        .reply(200, sampleMainnetTokenList)
+        .persist();
+
+      await controller.fetchTokenList(ChainId.mainnet);
+      await controller.fetchTokenList(ChainId.mainnet);
+      await controller.fetchTokenList(ChainId.mainnet);
+
+      // Verify getAllKeys was NOT called again after initialization
+      // (getItem may be called for other reasons, but getAllKeys is only used in load/migrate)
+      expect(getAllKeysCallCount).toBe(getAllKeysCallsAfterInit);
+
+      controller.destroy();
+    });
   });
 
   describe('deprecated methods', () => {

packages/assets-controllers/src/TokenListController.ts

@@ -123,6 +123,12 @@ export class TokenListController extends StaticIntervalPollingController<TokenLi
 > {
   readonly #mutex = new Mutex();
 
+  /**
+   * Promise that resolves when initialization (loading cache from storage) is complete.
+   * Methods that access the cache should await this before proceeding.
+   */
+  readonly #initializationPromise: Promise<void>;
+
   // Storage key prefix for per-chain files
   static readonly #storageKeyPrefix = 'tokensChainsCache';
 
@@ -191,18 +197,9 @@ export class TokenListController extends StaticIntervalPollingController<TokenLi
     this.updatePreventPollingOnNetworkRestart(preventPollingOnNetworkRestart);
     this.#abortController = new AbortController();
 
-    // Load cache from StorageService on initialization and handle migration
-    this.#loadCacheFromStorage()
-      .then(() => {
-        // Migrate existing cache from state to StorageService if needed
-        return this.#migrateStateToStorage();
-      })
-      .catch((error) => {
-        console.error(
-          'TokenListController: Failed to load cache from storage:',
-          error,
-        );
-      });
+    // Load cache from StorageService on initialization and handle migration.
+    // Store the promise so other methods can await it to avoid race conditions.
+    this.#initializationPromise = this.#initializeFromStorage();
 
     if (onNetworkStateChange) {
       // TODO: Either fix this lint violation or explain why it's necessary to ignore.
@@ -222,11 +219,36 @@ export class TokenListController extends StaticIntervalPollingController<TokenLi
     }
   }
 
+  /**
+   * Initialize the controller by loading cache from storage and running migration.
+   * This method acquires the mutex to prevent race conditions with fetchTokenList.
+   *
+   * @returns A promise that resolves when initialization is complete.
+   */
+  async #initializeFromStorage(): Promise<void> {
+    const releaseLock = await this.#mutex.acquire();
+    try {
+      await this.#loadCacheFromStorage();
+      await this.#migrateStateToStorage();
+    } catch (error) {
+      console.error(
+        'TokenListController: Failed to initialize from storage:',
+        error,
+      );
+    } finally {
+      releaseLock();
+    }
+  }
+
   /**
    * Load tokensChainsCache from StorageService into state.
    * Loads all cached chains from separate per-chain files in parallel.
    * Called during initialization to restore cached data.
    *
+   * Note: This method merges loaded data with existing state to avoid
+   * overwriting any fresh data that may have been fetched concurrently.
+   * Caller must hold the mutex.
+   *
    * @returns A promise that resolves when loading is complete.
    */
   async #loadCacheFromStorage(): Promise<void> {
@@ -278,10 +300,17 @@ export class TokenListController extends StaticIntervalPollingController<TokenLi
         }
       });
 
-      // Load into state (all chains available for TokenDetectionController)
+      // Merge loaded cache with existing state, preferring existing data
+      // (which may be fresher if fetched during initialization)
       if (Object.keys(loadedCache).length > 0) {
         this.update((state) => {
-          state.tokensChainsCache = loadedCache;
+          // Only load chains that don't already exist in state
+          // This prevents overwriting fresh API data with stale cached data
+          for (const [chainId, cacheData] of Object.entries(loadedCache)) {
+            if (!state.tokensChainsCache[chainId as Hex]) {
+              state.tokensChainsCache[chainId as Hex] = cacheData;
+            }
+          }
         });
       }
     } catch (error) {
@@ -497,6 +526,10 @@ export class TokenListController extends StaticIntervalPollingController<TokenLi
    * @param chainId - The chainId of the current chain triggering the fetch.
    */
   async fetchTokenList(chainId: Hex): Promise<void> {
+    // Wait for initialization to complete before fetching
+    // This ensures we have loaded any cached data from storage first
+    await this.#initializationPromise;
+
     const releaseLock = await this.#mutex.acquire();
     try {
       if (this.isCacheValid(chainId)) {
@@ -571,6 +604,9 @@ export class TokenListController extends StaticIntervalPollingController<TokenLi
    * This clears both state and all per-chain files in StorageService.
    */
   async clearingTokenListData(): Promise<void> {
+    // Wait for initialization to complete before clearing
+    await this.#initializationPromise;
+
     // Clear state
     this.update((state) => {
       state.tokensChainsCache = {};
@@ -593,13 +629,6 @@ export class TokenListController extends StaticIntervalPollingController<TokenLi
           this.messenger.call('StorageService:removeItem', name, key),
         ),
       );
-
-      // Also remove old single-file storage if it exists (cleanup)
-      await this.messenger.call(
-        'StorageService:removeItem',
-        name,
-        TokenListController.#storageKeyPrefix,
-      );
     } catch (error) {
       console.error(
         'TokenListController: Failed to clear cache from storage:',