feat: Allow mocking JSON-RPC implementations

packages/snaps-jest/src/helpers.test.tsx

@@ -993,6 +993,125 @@ describe('installSnap', () => {
       await close();
       await closeServer();
     });
+
+    it('mocks a JSON-RPC implementation', async () => {
+      jest.spyOn(console, 'log').mockImplementation();
+
+      const { snapId, close: closeServer } = await getMockServer({
+        sourceCode: `
+          module.exports.onRpcRequest = async () => {
+            return await ethereum.request({
+              method: 'foo',
+            });
+          };
+        `,
+        manifest: getSnapManifest({
+          initialPermissions: {
+            'endowment:ethereum-provider': {},
+          },
+        }),
+      });
+
+      const { request, close, mockJsonRpc } = await installSnap(snapId);
+      const { unmock } = mockJsonRpc(({ method }) => {
+        return `${method}_mocked`;
+      });
+
+      const response = await request({
+        method: 'foo',
+      });
+
+      expect(response).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            result: 'foo_mocked',
+          },
+        }),
+      );
+
+      unmock();
+
+      const unmockedResponse = await request({
+        method: 'foo',
+      });
+
+      expect(unmockedResponse).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            error: expect.objectContaining({
+              code: -32601,
+              message: 'The method "foo" does not exist / is not available.',
+            }),
+          },
+        }),
+      );
+
+      // `close` is deprecated because the Jest environment will automatically
+      // close the Snap when the test finishes. However, we still need to close
+      // the Snap in this test because it's run outside the Jest environment.
+      await close();
+      await closeServer();
+    });
+  });
+
+  describe('mockJsonRpcOnce', () => {
+    it('mocks a JSON-RPC method', async () => {
+      jest.spyOn(console, 'log').mockImplementation();
+
+      const { snapId, close: closeServer } = await getMockServer({
+        sourceCode: `
+          module.exports.onRpcRequest = async () => {
+            return await ethereum.request({
+              method: 'foo',
+            });
+          };
+        `,
+        manifest: getSnapManifest({
+          initialPermissions: {
+            'endowment:ethereum-provider': {},
+          },
+        }),
+      });
+
+      const { request, close, mockJsonRpcOnce } = await installSnap(snapId);
+      mockJsonRpcOnce({
+        method: 'foo',
+        result: 'mock',
+      });
+
+      const response = await request({
+        method: 'foo',
+      });
+
+      expect(response).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            result: 'mock',
+          },
+        }),
+      );
+
+      const unmockedResponse = await request({
+        method: 'foo',
+      });
+
+      expect(unmockedResponse).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            error: expect.objectContaining({
+              code: -32601,
+              message: 'The method "foo" does not exist / is not available.',
+            }),
+          },
+        }),
+      );
+
+      // `close` is deprecated because the Jest environment will automatically
+      // close the Snap when the test finishes. However, we still need to close
+      // the Snap in this test because it's run outside the Jest environment.
+      await close();
+      await closeServer();
+    });
   });
 });
 

packages/snaps-jest/src/helpers.ts

@@ -211,6 +211,7 @@ export async function installSnap<
     onProtocolRequest,
     onClientRequest,
     mockJsonRpc,
+    mockJsonRpcOnce,
     close,
   } = await getEnvironment().installSnap(...resolvedOptions);
   /* eslint-enable @typescript-eslint/unbound-method */
@@ -233,6 +234,7 @@ export async function installSnap<
     onProtocolRequest,
     onClientRequest,
     mockJsonRpc,
+    mockJsonRpcOnce,
     close: async () => {
       log('Closing execution service.');
       logInfo(

packages/snaps-simulation/src/helpers.test.tsx

@@ -875,5 +875,129 @@ describe('helpers', () => {
       await close();
       await closeServer();
     });
+
+    it('mocks a JSON-RPC implementation', async () => {
+      jest.spyOn(console, 'log').mockImplementation();
+
+      const { snapId, close: closeServer } = await getMockServer({
+        sourceCode: `
+          module.exports.onRpcRequest = async () => {
+            return await ethereum.request({
+              method: 'foo',
+            });
+          };
+        `,
+        manifest: getSnapManifest({
+          initialPermissions: {
+            'endowment:ethereum-provider': {},
+          },
+        }),
+      });
+
+      const { request, close, mockJsonRpc } = await installSnap(snapId);
+      const { unmock } = mockJsonRpc(({ method }) => {
+        return `${method}_mocked`;
+      });
+
+      const response = await request({
+        method: 'foo',
+      });
+
+      expect(response).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            result: 'foo_mocked',
+          },
+        }),
+      );
+
+      unmock();
+
+      const unmockedResponse = await request({
+        method: 'foo',
+      });
+
+      expect(unmockedResponse).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            error: expect.objectContaining({
+              code: -32601,
+              message: 'The method "foo" does not exist / is not available.',
+            }),
+          },
+        }),
+      );
+
+      // `close` is deprecated because the Jest environment will automatically
+      // close the Snap when the test finishes. However, we still need to close
+      // the Snap in this test because it's run outside the Jest environment.
+      await close();
+      await closeServer();
+    });
+  });
+
+  describe('mockJsonRpcOnce', () => {
+    it('mocks a JSON-RPC method once', async () => {
+      jest.spyOn(console, 'log').mockImplementation();
+
+      const { snapId, close: closeServer } = await getMockServer({
+        sourceCode: `
+          module.exports.onRpcRequest = async () => {
+            return await ethereum.request({
+              method: 'foo',
+            });
+          };
+        `,
+        manifest: getSnapManifest({
+          initialPermissions: {
+            'endowment:ethereum-provider': {},
+          },
+        }),
+      });
+
+      const { request, close, mockJsonRpcOnce } = await installSnap(snapId);
+      mockJsonRpcOnce({
+        method: 'foo_2',
+        result: 'invalid_mock',
+      });
+
+      mockJsonRpcOnce({
+        method: 'foo',
+        result: 'mock',
+      });
+
+      const response = await request({
+        method: 'foo',
+      });
+
+      expect(response).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            result: 'mock',
+          },
+        }),
+      );
+
+      const unmockedResponse = await request({
+        method: 'foo',
+      });
+
+      expect(unmockedResponse).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            error: expect.objectContaining({
+              code: -32601,
+              message: 'The method "foo" does not exist / is not available.',
+            }),
+          },
+        }),
+      );
+
+      // `close` is deprecated because the Jest environment will automatically
+      // close the Snap when the test finishes. However, we still need to close
+      // the Snap in this test because it's run outside the Jest environment.
+      await close();
+      await closeServer();
+    });
   });
 });

packages/snaps-simulation/src/helpers.ts

@@ -1,7 +1,8 @@
 import { HandlerType } from '@metamask/snaps-utils';
 import { create } from '@metamask/superstruct';
-import type { CaipChainId } from '@metamask/utils';
+import type { CaipChainId, JsonRpcRequest } from '@metamask/utils';
 import { createModuleLogger } from '@metamask/utils';
+import { nanoid } from '@reduxjs/toolkit';
 
 import { rootLogger } from './logger';
 import type { SimulationOptions } from './options';
@@ -221,6 +222,35 @@ export type SnapHelpers = {
     unmock(): void;
   };
 
+  /**
+   * Mock a JSON-RPC request once. This will cause the snap to respond with the
+   * specified response when a request with the specified method is sent.
+   *
+   * @param mock - The mock options.
+   * @param mock.method - The JSON-RPC request method.
+   * @param mock.result - The JSON-RPC response, which will be returned when a
+   * request with the specified method is sent.
+   * @example
+   * import { installSnap } from '@metamask/snaps-jest';
+   *
+   * // In the test
+   * const snap = await installSnap();
+   * snap.mockJsonRpcOnce({ method: 'eth_accounts', result: ['0x1234'] });
+   *
+   * // In the Snap
+   * const response =
+   *   await ethereum.request({ method: 'eth_accounts' }); // ['0x1234']
+   *
+   * const response2 =
+   *   await ethereum.request({ method: 'eth_accounts' }); // Default behavior
+   */
+  mockJsonRpcOnce(mock: JsonRpcMockOptions): {
+    /**
+     * Remove the mock.
+     */
+    unmock(): void;
+  };
+
   /**
    * Close the page running the snap. This is mainly useful for cleaning up
    * the test environment, and calling it is not strictly necessary.
@@ -318,6 +348,33 @@ export function getHelpers({
     });
   };
 
+  const mockJsonRpc = (mock: JsonRpcMockOptions, once: boolean) => {
+    log('Mocking JSON-RPC request %o.', mock);
+
+    const id = nanoid();
+
+    if (typeof mock === 'function') {
+      store.dispatch(addJsonRpcMock({ id, implementation: mock, once }));
+    } else {
+      const { method, result } = create(mock, JsonRpcMockOptionsStruct);
+      const implementation = (request: JsonRpcRequest) => {
+        if (request.method === method) {
+          return result;
+        }
+        return undefined;
+      };
+      store.dispatch(addJsonRpcMock({ id, implementation, once }));
+    }
+
+    return {
+      unmock() {
+        log('Unmocking JSON-RPC request %o.', mock);
+
+        store.dispatch(removeJsonRpcMock(id));
+      },
+    };
+  };
+
   return {
     // This can't be async because it returns a `SnapRequest`.
     // eslint-disable-next-line @typescript-eslint/promise-function-async
@@ -555,18 +612,11 @@ export function getHelpers({
     },
 
     mockJsonRpc(mock: JsonRpcMockOptions) {
-      log('Mocking JSON-RPC request %o.', mock);
-
-      const { method, result } = create(mock, JsonRpcMockOptionsStruct);
-      store.dispatch(addJsonRpcMock({ method, result }));
-
-      return {
-        unmock() {
-          log('Unmocking JSON-RPC request %o.', mock);
+      return mockJsonRpc(mock, false);
+    },
 
-          store.dispatch(removeJsonRpcMock(method));
-        },
-      };
+    mockJsonRpcOnce(mock: JsonRpcMockOptions) {
+      return mockJsonRpc(mock, true);
     },
 
     close: async () => {

packages/snaps-simulation/src/middleware/mock.test.ts

@@ -10,8 +10,9 @@ describe('createMockMiddleware', () => {
     const { store } = createStore(getMockOptions());
     store.dispatch(
       addJsonRpcMock({
-        method: 'foo',
-        result: 'bar',
+        id: 'foo',
+        implementation: () => 'bar',
+        once: false,
       }),
     );