From bc6864633064f7cbb04903dbd1e90feb4fadd031 Mon Sep 17 00:00:00 2001
From: Jongsun Suh <jongsun.suh@icloud.com>
Date: Fri, 16 Feb 2024 23:24:13 -0500
Subject: [PATCH] Replace `#getCorrectChainId` with
 `#getCorrectChainIdAndNetworkClientId`

---
 .../src/TokenDetectionController.ts           | 63 +++++++++++--------
 1 file changed, 37 insertions(+), 26 deletions(-)

diff --git a/packages/assets-controllers/src/TokenDetectionController.ts b/packages/assets-controllers/src/TokenDetectionController.ts
index fab50f1b8f..525864a5f8 100644
--- a/packages/assets-controllers/src/TokenDetectionController.ts
+++ b/packages/assets-controllers/src/TokenDetectionController.ts
@@ -234,7 +234,9 @@ export class TokenDetectionController extends StaticIntervalPollingController<
       selectedAddress ??
       this.messagingSystem.call('AccountsController:getSelectedAccount')
         .address;
-    this.#chainId = this.#getCorrectChainId(networkClientId);
+    const { chainId } =
+      this.#getCorrectChainIdAndNetworkClientId(networkClientId);
+    this.#chainId = chainId;
 
     const { useTokenDetection: defaultUseTokenDetection } =
       this.messagingSystem.call('PreferencesController:getState');
@@ -325,7 +327,8 @@ export class TokenDetectionController extends StaticIntervalPollingController<
         const isNetworkClientIdChanged =
           this.#networkClientId !== selectedNetworkClientId;
 
-        const newChainId = this.#getCorrectChainId(selectedNetworkClientId);
+        const { chainId: newChainId } =
+          this.#getCorrectChainIdAndNetworkClientId(selectedNetworkClientId);
         this.#isDetectionEnabledForNetwork =
           isTokenDetectionSupportedForNetwork(newChainId);
 
@@ -397,13 +400,33 @@ export class TokenDetectionController extends StaticIntervalPollingController<
     }, this.getIntervalLength());
   }
 
-  #getCorrectChainId(networkClientId?: NetworkClientId) {
-    const { chainId } =
-      this.messagingSystem.call(
+  #getCorrectChainIdAndNetworkClientId(networkClientId?: NetworkClientId): {
+    chainId: Hex;
+    networkClientId: NetworkClientId;
+  } {
+    if (networkClientId) {
+      const networkConfiguration = this.messagingSystem.call(
         'NetworkController:getNetworkConfigurationByNetworkClientId',
-        networkClientId ?? this.#networkClientId,
-      ) ?? {};
-    return chainId ?? this.#chainId;
+        networkClientId,
+      );
+      if (networkConfiguration) {
+        return {
+          chainId: networkConfiguration.chainId,
+          networkClientId,
+        };
+      }
+    }
+    const { chainId } = this.messagingSystem.call(
+      'NetworkController:getProviderConfig',
+    );
+    const newNetworkClientId = this.messagingSystem.call(
+      'NetworkController:findNetworkClientIdByChainId',
+      this.#chainId,
+    );
+    return {
+      chainId,
+      networkClientId: newNetworkClientId,
+    };
   }
 
   async _executePoll(
@@ -459,25 +482,13 @@ export class TokenDetectionController extends StaticIntervalPollingController<
     if (!this.isActive) {
       return;
     }
+
     const addressAgainstWhichToDetect = accountAddress ?? this.#selectedAddress;
-    let chainIdAgainstWhichToDetect: Hex;
-    let networkClientIdAgainstWhichToDetect: NetworkClientId;
-    if (networkClientId) {
-      const networkConfiguration = this.messagingSystem.call(
-        'NetworkController:getNetworkConfigurationByNetworkClientId',
-        networkClientId,
-      );
-      chainIdAgainstWhichToDetect = networkConfiguration.chainId;
-      networkClientIdAgainstWhichToDetect = networkClientId;
-    } else {
-      chainIdAgainstWhichToDetect = this.messagingSystem.call(
-        'NetworkController:getProviderConfig',
-      ).chainId;
-      networkClientIdAgainstWhichToDetect = this.messagingSystem.call(
-        'NetworkController:findNetworkClientIdByChainId',
-        chainIdAgainstWhichToDetect,
-      );
-    }
+    const {
+      chainId: chainIdAgainstWhichToDetect,
+      networkClientId: networkClientIdAgainstWhichToDetect,
+    } = this.#getCorrectChainIdAndNetworkClientId(networkClientId);
+
     if (!isTokenDetectionSupportedForNetwork(chainIdAgainstWhichToDetect)) {
       return;
     }