Add snap metadata to account if owned

GuillaumeRx · GuillaumeRx · commit 2f6a404f1a73 · 2025-07-03T16:57:50.000+02:00
diff --git a/packages/snaps-simulation/src/constants.ts b/packages/snaps-simulation/src/constants.ts
@@ -40,6 +40,7 @@ export const DEFAULT_ACCOUNTS: SimulationAccount[] = [
     id: '29bc7513-d1b9-4466-98a6-f5f9e0b90137',
     scopes: ['eip155:0'],
     selected: false,
+    owned: false,
     // We don't expose assets for EVM accounts as it's not supported in the AssetSelector.
     assets: [],
   },
@@ -52,6 +53,7 @@ export const DEFAULT_ACCOUNTS: SimulationAccount[] = [
       'solana:4uhcVJyU9pJkvQyS88uRDiswHXSCkY3z',
     ],
     selected: true,
+    owned: true,
     assets: [
       'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501',
       'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v',
diff --git a/packages/snaps-simulation/src/options.test.ts b/packages/snaps-simulation/src/options.test.ts
@@ -11,6 +11,7 @@ describe('getOptions', () => {
             "address": "0x1234567890abcdef1234567890abcdef12345678",
             "assets": [],
             "id": "29bc7513-d1b9-4466-98a6-f5f9e0b90137",
+            "owned": false,
             "scopes": [
               "eip155:0",
             ],
@@ -23,6 +24,7 @@ describe('getOptions', () => {
               "solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
             ],
             "id": "e051723c-85d0-43a3-b9bf-568a90d3f378",
+            "owned": true,
             "scopes": [
               "solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp",
               "solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1",
@@ -72,6 +74,7 @@ describe('getOptions', () => {
             "address": "0x1234567890abcdef1234567890abcdef12345678",
             "assets": [],
             "id": "29bc7513-d1b9-4466-98a6-f5f9e0b90137",
+            "owned": false,
             "scopes": [
               "eip155:0",
             ],
@@ -84,6 +87,7 @@ describe('getOptions', () => {
               "solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
             ],
             "id": "e051723c-85d0-43a3-b9bf-568a90d3f378",
+            "owned": true,
             "scopes": [
               "solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp",
               "solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1",
diff --git a/packages/snaps-simulation/src/options.ts b/packages/snaps-simulation/src/options.ts
@@ -29,6 +29,7 @@ const SimulationAccountStruct = object({
   id: string(),
   scopes: array(CaipChainIdStruct),
   selected: defaulted(optional(boolean()), false),
+  owned: defaulted(optional(boolean()), false),
   assets: defaulted(optional(array(CaipAssetTypeStruct)), []),
 });
 
diff --git a/packages/snaps-simulation/src/simulation.test.ts b/packages/snaps-simulation/src/simulation.test.ts
@@ -10,7 +10,10 @@ import { DIALOG_APPROVAL_TYPES } from '@metamask/snaps-rpc-methods';
 import type { CaipAssetType, CaipChainId } from '@metamask/snaps-sdk';
 import { AuxiliaryFileEncoding, text } from '@metamask/snaps-sdk';
 import { VirtualFile } from '@metamask/snaps-utils';
-import { getSnapManifest } from '@metamask/snaps-utils/test-utils';
+import {
+  getSnapManifest,
+  MOCK_SNAP_ID,
+} from '@metamask/snaps-utils/test-utils';
 
 import { DEFAULT_SRP } from './constants';
 import {
@@ -26,6 +29,7 @@ import {
   getRestrictedSnapInterfaceControllerMessenger,
   getRootControllerMessenger,
 } from './test-utils';
+import { addSnapMetadataToAccount } from './utils/account';
 
 describe('installSnap', () => {
   it('installs a Snap and returns the execution service', async () => {
@@ -693,15 +697,15 @@ describe('registerActions', () => {
   const controllerMessenger = getRootControllerMessenger(false);
 
   it('registers `PhishingController:testOrigin`', async () => {
-    registerActions(controllerMessenger, runSaga, options);
+    registerActions(controllerMessenger, runSaga, options, MOCK_SNAP_ID);
 
     expect(
       controllerMessenger.call('PhishingController:testOrigin', 'foo'),
     ).toStrictEqual({ result: false, type: 'all' });
   });
 
   it('registers `ApprovalController:hasRequest`', async () => {
-    registerActions(controllerMessenger, runSaga, options);
+    registerActions(controllerMessenger, runSaga, options, MOCK_SNAP_ID);
 
     store.dispatch(
       setInterface({ type: DIALOG_APPROVAL_TYPES.default, id: 'foo' }),
@@ -713,7 +717,7 @@ describe('registerActions', () => {
   });
 
   it('registers `ApprovalController:acceptRequest`', async () => {
-    registerActions(controllerMessenger, runSaga, options);
+    registerActions(controllerMessenger, runSaga, options, MOCK_SNAP_ID);
 
     store.dispatch(
       setInterface({ type: DIALOG_APPROVAL_TYPES.default, id: 'foo' }),
@@ -729,36 +733,59 @@ describe('registerActions', () => {
   });
 
   it('registers `AccountsController:getAccountByAddress`', async () => {
-    registerActions(controllerMessenger, runSaga, options);
+    registerActions(controllerMessenger, runSaga, options, MOCK_SNAP_ID);
 
     expect(
       controllerMessenger.call(
         'AccountsController:getAccountByAddress',
         mockedAccounts[0].address,
       ),
-    ).toStrictEqual(mockedAccounts[0]);
+    ).toStrictEqual(addSnapMetadataToAccount(mockedAccounts[0], MOCK_SNAP_ID));
+
+    expect(
+      controllerMessenger.call('AccountsController:getAccountByAddress', 'foo'),
+    ).toBeUndefined();
   });
 
   it('registers `AccountsController:getSelectedMultichainAccount`', async () => {
-    registerActions(controllerMessenger, runSaga, options);
+    registerActions(controllerMessenger, runSaga, options, MOCK_SNAP_ID);
 
     expect(
       controllerMessenger.call(
         'AccountsController:getSelectedMultichainAccount',
       ),
-    ).toStrictEqual(mockedAccounts[0]);
+    ).toStrictEqual(addSnapMetadataToAccount(mockedAccounts[0], MOCK_SNAP_ID));
+  });
+
+  it('returns `undefined` in `AccountsController:getSelectedMultichainAccount` if no account is selected', async () => {
+    registerActions(
+      controllerMessenger,
+      runSaga,
+      { ...options, accounts: [] },
+      MOCK_SNAP_ID,
+    );
+
+    expect(
+      controllerMessenger.call(
+        'AccountsController:getSelectedMultichainAccount',
+      ),
+    ).toBeUndefined();
   });
 
   it('registers `AccountsController:listMultichainAccounts`', async () => {
-    registerActions(controllerMessenger, runSaga, options);
+    registerActions(controllerMessenger, runSaga, options, MOCK_SNAP_ID);
 
     expect(
       controllerMessenger.call('AccountsController:listMultichainAccounts'),
-    ).toStrictEqual(mockedAccounts);
+    ).toStrictEqual(
+      mockedAccounts.map((account) =>
+        addSnapMetadataToAccount(account, MOCK_SNAP_ID),
+      ),
+    );
   });
 
   it('registers `MultichainAssetsController:getState`', async () => {
-    registerActions(controllerMessenger, runSaga, options);
+    registerActions(controllerMessenger, runSaga, options, MOCK_SNAP_ID);
 
     expect(
       controllerMessenger.call('MultichainAssetsController:getState'),
diff --git a/packages/snaps-simulation/src/simulation.ts b/packages/snaps-simulation/src/simulation.ts
@@ -50,6 +50,7 @@ import type { SimulationOptions, SimulationUserOptions } from './options';
 import { getOptions } from './options';
 import type { Interface, RunSagaFunction, Store } from './store';
 import { createStore, getCurrentInterface } from './store';
+import { addSnapMetadataToAccount } from './utils/account';
 
 /**
  * Options for the execution service, without the options that are shared
@@ -303,7 +304,7 @@ export async function installSnap<
 
   const controllerMessenger = new Messenger<any, any>();
 
-  registerActions(controllerMessenger, runSaga, options);
+  registerActions(controllerMessenger, runSaga, options, snapId);
 
   // Set up controllers and JSON-RPC stack.
   const restrictedHooks = getRestrictedHooks(options);
@@ -474,11 +475,13 @@ export function getPermittedHooks(
  * @param controllerMessenger - The controller messenger.
  * @param runSaga - The run saga function.
  * @param options - The simulation options.
+ * @param snapId - The ID of the Snap.
  */
 export function registerActions(
   controllerMessenger: RootControllerMessenger,
   runSaga: RunSagaFunction,
   options: SimulationOptions,
+  snapId: SnapId,
 ) {
   controllerMessenger.registerActionHandler(
     'PhishingController:testOrigin',
@@ -487,24 +490,47 @@ export function registerActions(
 
   controllerMessenger.registerActionHandler(
     'AccountsController:getAccountByAddress',
-    (address) =>
-      // @ts-expect-error - This is a partial account with only the necessary
-      // data used by the interface controller.
-      options.accounts.find((account) => address === account.address),
+    // @ts-expect-error - This is a partial account with only the necessary
+    // data used by the interface controller.
+    (address) => {
+      const matchingAccount = options.accounts.find(
+        (account) => address === account.address,
+      );
+
+      if (!matchingAccount) {
+        return undefined;
+      }
+
+      return addSnapMetadataToAccount(matchingAccount, snapId);
+    },
   );
 
   controllerMessenger.registerActionHandler(
     'AccountsController:getSelectedMultichainAccount',
     // @ts-expect-error - This is a partial account with only the necessary
     // data used by the interface controller.
-    () => options.accounts.find((account) => account.selected),
+    () => {
+      const selectedAccount = options.accounts.find(
+        (account) => account.selected,
+      );
+
+      if (!selectedAccount) {
+        return undefined;
+      }
+
+      return addSnapMetadataToAccount(selectedAccount, snapId);
+    },
   );
 
   controllerMessenger.registerActionHandler(
     'AccountsController:listMultichainAccounts',
-    // @ts-expect-error - These are partial accounts with only the necessary
-    // data used by the interface controller.
-    () => options.accounts,
+
+    () =>
+      // @ts-expect-error - These are partial accounts with only the necessary
+      // data used by the interface controller.
+      options.accounts.map((account) =>
+        addSnapMetadataToAccount(account, snapId),
+      ),
   );
 
   controllerMessenger.registerActionHandler(
diff --git a/packages/snaps-simulation/src/utils/account.test.ts b/packages/snaps-simulation/src/utils/account.test.ts
@@ -0,0 +1,52 @@
+import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils';
+import type { SimulationAccount } from 'src/options';
+
+import { addSnapMetadataToAccount } from './account';
+
+describe('addSnapMetadataToAccount', () => {
+  it('adds snap metadata to an owned account', () => {
+    const account: SimulationAccount = {
+      id: 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
+      selected: true,
+      address: '7S3P4HxJpyyigGzodYwHtCxZyUQe9JiBMHyRWXArAaKv',
+      owned: true,
+      scopes: ['solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp'],
+      assets: [
+        'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:105',
+        'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v',
+      ],
+    };
+
+    const result = addSnapMetadataToAccount(account, MOCK_SNAP_ID);
+
+    expect(result).toStrictEqual({
+      ...account,
+      metadata: {
+        snap: {
+          id: MOCK_SNAP_ID,
+        },
+      },
+    });
+  });
+
+  it('returns an account without snap metadata if not owned', () => {
+    const account: SimulationAccount = {
+      id: 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
+      selected: true,
+      address: '7S3P4HxJpyyigGzodYwHtCxZyUQe9JiBMHyRWXArAaKv',
+      owned: false,
+      scopes: ['solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp'],
+      assets: [
+        'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:105',
+        'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v',
+      ],
+    };
+
+    const result = addSnapMetadataToAccount(account, MOCK_SNAP_ID);
+
+    expect(result).toStrictEqual({
+      ...account,
+      metadata: {},
+    });
+  });
+});
diff --git a/packages/snaps-simulation/src/utils/account.ts b/packages/snaps-simulation/src/utils/account.ts
@@ -0,0 +1,23 @@
+import type { SnapId } from '@metamask/snaps-sdk';
+import type { SimulationAccount } from 'src/options';
+
+export const addSnapMetadataToAccount = (
+  account: SimulationAccount,
+  snapId: SnapId,
+) => {
+  if (!account.owned) {
+    return {
+      ...account,
+      metadata: {},
+    };
+  }
+
+  return {
+    ...account,
+    metadata: {
+      snap: {
+        id: snapId,
+      },
+    },
+  };
+};
