diff --git a/packages/snap/snap.manifest.json b/packages/snap/snap.manifest.json
index abea7f3..6a45358 100644
--- a/packages/snap/snap.manifest.json
+++ b/packages/snap/snap.manifest.json
@@ -7,7 +7,7 @@
"url": "git+https://github.com/MetaMask/snap-institutional-wallet.git"
},
"source": {
- "shasum": "2kLOK4/mWAaSYQkQZ5ABdIUtlrTQXgnrzZWFudNjN9s=",
+ "shasum": "xC69NMxzQ9d+X0HrsOyJJM6Pna8zpxzeSdt3WVYySaw=",
"location": {
"npm": {
"filePath": "dist/bundle.js",
diff --git a/packages/snap/src/index.ts b/packages/snap/src/index.ts
index b2bf830..7f167eb 100644
--- a/packages/snap/src/index.ts
+++ b/packages/snap/src/index.ts
@@ -34,7 +34,6 @@ import type { SnapContext } from './lib/types/Context';
import { CustodianApiMap, CustodianType } from './lib/types/CustodianType';
import logger from './logger';
import { InternalMethod, originPermissions } from './permissions';
-// @audit - this file needs unittests
/**
* Verify if the caller can call the requested method.
diff --git a/packages/snap/src/keyring.test.ts b/packages/snap/src/keyring.test.ts
new file mode 100644
index 0000000..523c396
--- /dev/null
+++ b/packages/snap/src/keyring.test.ts
@@ -0,0 +1,330 @@
+import { MethodNotFoundError } from '@metamask/snaps-sdk';
+
+import { CustodialKeyring } from './keyring';
+import { TOKEN_EXPIRED_EVENT } from './lib/custodian-types/constants';
+import { CustodianApiMap } from './lib/types/CustodianType';
+import type { ICustodianApi } from './lib/types/ICustodianApi';
+
+// Mock dependencies
+jest.mock('./lib/custodian-types/custodianMetadata', () => ({
+ custodianMetadata: {
+ BitGo: {
+ enabled: false,
+ apiBaseUrl: 'https://mock-url.com',
+ apiVersion: 'BitGo',
+ custodianPublishesTransaction: true,
+ iconUrl: 'https://mock-url.com/icon.svg',
+ },
+ ECA3: {
+ enabled: true,
+ apiBaseUrl: 'https://mock-url.com',
+ apiVersion: 'ECA3',
+ custodianPublishesTransaction: false,
+ iconUrl: 'https://mock-url.com/icon.svg',
+ },
+ },
+}));
+
+jest.mock('./features/info-message/rendex');
+jest.mock('@metamask/keyring-api', () => ({
+ ...jest.requireActual('@metamask/keyring-api'),
+ emitSnapKeyringEvent: jest.fn(),
+}));
+
+jest.mock('./lib/types/CustodianType', () => ({
+ CustodianType: {
+ ECA3: 'ECA3',
+ ECA1: 'ECA1',
+ BitGo: 'BitGo',
+ Cactus: 'Cactus',
+ },
+ CustodianApiMap: {
+ ECA3: jest.fn(),
+ },
+}));
+
+describe('CustodialKeyring', () => {
+ let keyring: CustodialKeyring;
+ let mockStateManager: any;
+ let mockRequestManager: any;
+
+ beforeEach(() => {
+ // Reset all mocks
+ jest.clearAllMocks();
+
+ // Setup mock state manager
+ mockStateManager = {
+ listAccounts: jest.fn(),
+ getAccount: jest.fn(),
+ withTransaction: jest.fn(),
+ removeAccounts: jest.fn(),
+ getWalletByAddress: jest.fn(),
+ listWallets: jest.fn(),
+ updateWalletDetails: jest.fn(),
+ addWallet: jest.fn(),
+ };
+
+ // Setup mock request manager
+ mockRequestManager = {
+ upsertRequest: jest.fn(),
+ listRequests: jest.fn(),
+ };
+
+ // Create keyring instance
+ keyring = new CustodialKeyring(mockStateManager, mockRequestManager);
+ });
+
+ describe('listAccounts', () => {
+ it('should return accounts from state manager', async () => {
+ const mockAccounts = [{ id: '1', address: '0x123' }];
+ mockStateManager.listAccounts.mockResolvedValue(mockAccounts);
+
+ const result = await keyring.listAccounts();
+ expect(result).toStrictEqual(mockAccounts);
+ expect(mockStateManager.listAccounts).toHaveBeenCalled();
+ });
+ });
+
+ describe('getAccount', () => {
+ it('should return account from state manager', async () => {
+ const mockAccount = { id: '1', address: '0x123' };
+ mockStateManager.getAccount.mockResolvedValue(mockAccount);
+
+ const result = await keyring.getAccount('1');
+ expect(result).toStrictEqual(mockAccount);
+ expect(mockStateManager.getAccount).toHaveBeenCalledWith('1');
+ });
+
+ it('should return undefined for non-existent account', async () => {
+ mockStateManager.getAccount.mockResolvedValue(null);
+
+ const result = await keyring.getAccount('nonexistent');
+ expect(result).toBeUndefined();
+ });
+ });
+
+ describe('filterAccountChains', () => {
+ it('should filter supported chains', async () => {
+ const mockAccount = { id: '1', address: '0x123' };
+ const mockWallet = {
+ account: mockAccount,
+ details: {
+ custodianType: 'ECA3',
+ token: 'mock-token',
+ custodianApiUrl: 'https://api.example.com',
+ refreshTokenUrl: 'https://refresh.example.com',
+ },
+ };
+
+ mockStateManager.getAccount.mockResolvedValue(mockAccount);
+ mockStateManager.getWalletByAddress.mockResolvedValue(mockWallet);
+
+ // Mock the custodian API to return decimal chain IDs
+ const mockCustodianApi = {
+ getSupportedChains: jest.fn().mockResolvedValue(['eip155:1']), // CAIP-2 format
+ };
+
+ jest
+ .spyOn(keyring as any, 'getCustodianApiForAddress')
+ .mockImplementation(async () => Promise.resolve(mockCustodianApi));
+
+ const result = await keyring.filterAccountChains('1', ['0x1', '0x2']);
+ expect(result).toStrictEqual(['0x1']);
+ });
+
+ it('should throw error for non-existent account', async () => {
+ mockStateManager.getAccount.mockResolvedValue(null);
+
+ await expect(keyring.filterAccountChains('1', [])).rejects.toThrow(
+ "Account '1' not found",
+ );
+ });
+ });
+
+ describe('updateAccount', () => {
+ it('should throw MethodNotFoundError', async () => {
+ await expect(keyring.updateAccount({} as any)).rejects.toThrow(
+ MethodNotFoundError,
+ );
+ });
+ });
+
+ describe('deleteAccount', () => {
+ it('should delete account and emit event', async () => {
+ mockStateManager.withTransaction.mockImplementation((callback: any) =>
+ callback(),
+ );
+
+ await keyring.deleteAccount('1');
+
+ expect(mockStateManager.removeAccounts).toHaveBeenCalledWith(['1']);
+ expect(mockStateManager.withTransaction).toHaveBeenCalled();
+ });
+ });
+
+ describe('submitRequest', () => {
+ it('should handle personal sign request', async () => {
+ const mockAccount = {
+ id: '1',
+ address: '0x123',
+ options: {
+ custodian: {
+ displayName: 'Test Custodian',
+ deferPublication: false,
+ },
+ },
+ };
+ const mockRequest = {
+ id: 'request-1',
+ scope: 'scope-1',
+ account: mockAccount.id,
+ request: {
+ method: 'personal_sign',
+ params: ['message', mockAccount.address],
+ },
+ };
+
+ const mockWallet = {
+ account: mockAccount,
+ details: {
+ custodianType: 'ECA3',
+ token: 'mock-token',
+ custodianApiUrl: 'https://api.example.com',
+ refreshTokenUrl: 'https://refresh.example.com',
+ },
+ };
+
+ mockStateManager.getAccount.mockResolvedValue(mockAccount);
+ mockStateManager.getWalletByAddress.mockResolvedValue(mockWallet);
+
+ // Mock the custodian API
+ const mockCustodianApi = {
+ signPersonalMessage: jest.fn().mockResolvedValue({ id: 'msg-1' }),
+ getSignedMessageLink: jest.fn().mockResolvedValue({
+ text: 'View Message',
+ id: 'msg-1',
+ url: 'https://example.com',
+ action: 'view',
+ }),
+ };
+ jest
+ .spyOn(keyring as any, 'getCustodianApiForAddress')
+ .mockResolvedValue(mockCustodianApi);
+
+ const result = await keyring.submitRequest(mockRequest);
+
+ expect(result).toStrictEqual({ pending: true });
+ expect(mockRequestManager.upsertRequest).toHaveBeenCalledWith(
+ expect.objectContaining({
+ type: 'message',
+ subType: 'personalSign',
+ }),
+ );
+ });
+
+ // Add more tests for other request types...
+ });
+
+ describe('getCustodianApiForAddress', () => {
+ it('should handle token expiry events and update wallet details', async () => {
+ const mockAddress = '0x123';
+ const mockWallets = [
+ {
+ account: { id: '1', address: mockAddress },
+ details: {
+ token: 'oldToken',
+ custodianApiUrl: 'https://api.example.com',
+ custodianType: 'ECA3',
+ refreshTokenUrl: 'https://refresh.example.com',
+ },
+ },
+ {
+ account: { id: '2', address: '0x456' },
+ details: {
+ token: 'oldToken',
+ custodianApiUrl: 'https://api.example.com',
+ custodianType: 'ECA3',
+ refreshTokenUrl: 'https://refresh.example.com',
+ },
+ },
+ {
+ account: { id: '3', address: '0x789' },
+ details: {
+ token: 'differentToken',
+ custodianApiUrl: 'https://different.api.com',
+ custodianType: 'ECA3',
+ refreshTokenUrl: 'https://refresh.example.com',
+ },
+ },
+ ];
+
+ const updatedDetails = new Map();
+
+ mockStateManager.getWalletByAddress.mockImplementation(
+ async (address: string) => {
+ return Promise.resolve(
+ mockWallets.find((wallet) => wallet.account.address === address),
+ );
+ },
+ );
+
+ mockStateManager.listWallets.mockImplementation(async () => {
+ return Promise.resolve(mockWallets);
+ });
+
+ mockStateManager.updateWalletDetails.mockImplementation(
+ async (id: string, details: any) => {
+ updatedDetails.set(id, { ...details });
+ return Promise.resolve();
+ },
+ );
+
+ let tokenEventCallback: ((event: any) => Promise<void>) | undefined;
+ const mockCustodianApi: Partial<ICustodianApi> = {
+ on: jest.fn().mockImplementation((event, callback) => {
+ if (event === TOKEN_EXPIRED_EVENT) {
+ tokenEventCallback = callback;
+ }
+ }),
+ getSupportedChains: jest.fn(),
+ };
+
+ const ECA3Mock =
+ CustodianApiMap.ECA3 as unknown as jest.Mock<ICustodianApi>;
+ ECA3Mock.mockImplementation(() => mockCustodianApi as ICustodianApi);
+
+ // Call the method under test
+ const api = await (keyring as any).getCustodianApiForAddress(mockAddress);
+
+ // Simulate token expiry event and wait for all promises to resolve
+ if (tokenEventCallback) {
+ await tokenEventCallback({
+ oldRefreshToken: 'oldToken',
+ newRefreshToken: 'newToken',
+ apiUrl: 'https://api.example.com',
+ });
+ }
+
+ // Wait for any pending promises
+ await new Promise((resolve) => setTimeout(resolve, 0));
+
+ // Verify each wallet update
+ const expectedDetails = {
+ token: 'newToken',
+ custodianApiUrl: 'https://api.example.com',
+ custodianType: 'ECA3',
+ refreshTokenUrl: 'https://refresh.example.com',
+ };
+
+ expect(updatedDetails.get('1')).toStrictEqual(expectedDetails);
+ expect(updatedDetails.get('2')).toStrictEqual(expectedDetails);
+ expect(updatedDetails.has('3')).toBe(false);
+
+ // Verify the event listener was set up
+ expect(api.on).toHaveBeenCalledWith(
+ TOKEN_EXPIRED_EVENT,
+ expect.any(Function),
+ );
+ });
+ });
+});
diff --git a/packages/snap/src/keyring.tsx b/packages/snap/src/keyring.tsx
index 81b5e66..2500564 100644
--- a/packages/snap/src/keyring.tsx
+++ b/packages/snap/src/keyring.tsx
@@ -52,7 +52,7 @@ type RequestManagerFacade = {
};
export class CustodialKeyring implements Keyring {
- #custodianApi: Record<string, ICustodianApi>; // maps address to memoized custodian api
+ #custodianApi: Map<string, ICustodianApi>;
#requestManagerFacade: RequestManagerFacade;
@@ -63,7 +63,7 @@ export class CustodialKeyring implements Keyring {
requestManagerFacade: RequestManagerFacade,
) {
this.#stateManager = stateManager;
- this.#custodianApi = {};
+ this.#custodianApi = new Map<string, ICustodianApi>();
this.#requestManagerFacade = requestManagerFacade;
}
@@ -221,13 +221,13 @@ export class CustodialKeyring implements Keyring {
}
async getCustodianApiForAddress(address: string): Promise<ICustodianApi> {
- if (!this.#custodianApi[address]) {
+ if (!this.#custodianApi.has(address)) {
const wallet = await this.#stateManager.getWalletByAddress(address);
if (!wallet) {
throw new Error(`Wallet for account ${address} does not exist`);
}
const custodianApi = this.#getCustodianApi(wallet.details);
- this.#custodianApi[address] = custodianApi;
+ this.#custodianApi.set(address, custodianApi);
custodianApi.on(
TOKEN_EXPIRED_EVENT,
(payload: IRefreshTokenChangeEvent) => {
@@ -235,7 +235,7 @@ export class CustodialKeyring implements Keyring {
},
);
}
- return this.#custodianApi[address] as ICustodianApi;
+ return this.#custodianApi.get(address) as ICustodianApi;
}
#getCustodianApi(details: OnBoardingRpcRequest): ICustodianApi {
@@ -259,20 +259,21 @@ export class CustodialKeyring implements Keyring {
wallet.details.token === payload.oldRefreshToken &&
wallet.details.custodianApiUrl === payload.apiUrl,
);
- // Update the custodian api for each wallet
- walletsToUpdate.forEach((wallet) => {
- wallet.details.token = payload.newRefreshToken;
- });
-
- // Delete the custodian api from the cache for each address
- walletsToUpdate.forEach((wallet) => {
- delete this.#custodianApi[wallet.account.address];
- });
for (const wallet of walletsToUpdate) {
+ // Create new details object with updated token
+ const updatedDetails = {
+ ...wallet.details,
+ token: payload.newRefreshToken,
+ };
+
+ // Clear cache
+ this.#custodianApi.delete(wallet.account.address);
+
+ // Update state with new details
await this.#stateManager.updateWalletDetails(
wallet.account.id,
- wallet.details,
+ updatedDetails,
);
}
}