From 8d1b9b33f5ddb1e617d059c0aa32ea53292d5eb3 Mon Sep 17 00:00:00 2001
From: Christopher Ferreira <christopher.ferreira@consensys.net>
Date: Wed, 8 Jan 2025 14:17:58 +0000
Subject: [PATCH] chore: uses await instead of .then and adds coverage

---
 .../setupExtensionPreferences.test.ts         | 34 ++++++++++++++++
 .../setupExtensionPreferences.ts              | 39 +++++++------------
 2 files changed, 49 insertions(+), 24 deletions(-)

diff --git a/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.test.ts b/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.test.ts
index 8941b891b..3b13e4f2c 100644
--- a/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.test.ts
+++ b/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.test.ts
@@ -1,6 +1,8 @@
 import { TrackingEvents } from '@metamask/sdk-communication-layer';
 import { MetaMaskSDK } from '../../../sdk';
+import { SDKProvider } from '../../../provider/SDKProvider';
 import { getBrowserExtension } from '../../../utils/get-browser-extension';
+import { EXTENSION_EVENTS, RPC_METHODS } from '../../../config';
 import { setupExtensionPreferences } from './setupExtensionPreferences';
 
 jest.mock('../../../utils/get-browser-extension');
@@ -81,4 +83,36 @@ describe('setupExtensionPreferences', () => {
       event: TrackingEvents.SDK_USE_EXTENSION,
     });
   });
+
+  it('should terminate instance if permissions are revoked', async () => {
+    const mockTerminate = jest.fn().mockResolvedValue(undefined);
+    const mockRequest = jest.fn().mockResolvedValue([]);
+
+    instance.getProvider = jest.fn().mockReturnValue({
+      request: mockRequest,
+    }) as unknown as () => SDKProvider;
+
+    instance.terminate = mockTerminate;
+    instance.extensionActive = true; // Ensure extension is active
+
+    const mockBrowserExtension = {
+      on: jest.fn((event: string, callback: (accounts: string[]) => void) => {
+        if (event === EXTENSION_EVENTS.ACCOUNTS_CHANGED) {
+          // Directly call the callback to simulate the event
+          callback([]);
+        }
+      }),
+    };
+
+    (getBrowserExtension as jest.Mock).mockReturnValue(mockBrowserExtension);
+
+    await setupExtensionPreferences(instance);
+
+    expect(mockRequest).toHaveBeenCalledWith({
+      method: RPC_METHODS.WALLET_GETPERMISSIONS,
+      params: [],
+    });
+
+    expect(mockTerminate).toHaveBeenCalled();
+  });
 });
diff --git a/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.ts b/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.ts
index ee137cc42..4cfb46499 100644
--- a/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.ts
+++ b/packages/sdk/src/services/MetaMaskSDK/InitializerManager/setupExtensionPreferences.ts
@@ -68,7 +68,7 @@ export async function setupExtensionPreferences(instance: MetaMaskSDK) {
 
       metamaskBrowserExtension.on(
         EXTENSION_EVENTS.ACCOUNTS_CHANGED,
-        (accounts) => {
+        async (accounts) => {
           logger(
             `[MetaMaskSDK: setupExtensionPreferences()] PROPAGATE accountsChanged accounts=${accounts}`,
           );
@@ -84,37 +84,28 @@ export async function setupExtensionPreferences(instance: MetaMaskSDK) {
           }
 
           if (isExtensionActive && (accounts as string[])?.length === 0) {
-            instance
+            const getPermissionsResponse = await instance
               .getProvider()
               ?.request({
                 method: RPC_METHODS.WALLET_GETPERMISSIONS,
                 params: [],
-              })
-              .then((getPermissionsResponse) => {
-                const permissions = getPermissionsResponse as {
-                  caveats: { type: string; value: string[] }[];
-                  parentCapability: string;
-                }[];
-                logger(
-                  `[MetaMaskSDK: setupExtensionPreferences()] permissions`,
-                  permissions,
-                );
+              });
+
+            const permissions = getPermissionsResponse as {
+              caveats: { type: string; value: string[] }[];
+              parentCapability: string;
+            }[];
 
-                if (permissions.length === 0) {
-                  instance.terminate().catch((error) => {
-                    logger(
-                      `[MetaMaskSDK: setupExtensionPreferences()] error terminating on permissions revoked`,
-                      error,
-                    );
-                  });
-                }
-              })
-              .catch((error) => {
+            if (permissions.length === 0) {
+              try {
+                await instance.terminate();
+              } catch (error) {
                 logger(
-                  `[MetaMaskSDK: setupExtensionPreferences()] error getting permissions`,
+                  `[MetaMaskSDK: setupExtensionPreferences()] error terminating on permissions revoked`,
                   error,
                 );
-              });
+              }
+            }
           }
         },
       );