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,205 @@ 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();
+    });
+
+    it('supports queueing JSON-RPC mocks', 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',
+      });
+
+      mockJsonRpcOnce({
+        method: 'foo',
+        result: 'mock2',
+      });
+
+      const response1 = await request({
+        method: 'foo',
+      });
+
+      expect(response1).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            result: 'mock',
+          },
+        }),
+      );
+
+      const response2 = await request({
+        method: 'foo',
+      });
+
+      expect(response2).toStrictEqual(
+        expect.objectContaining({
+          response: {
+            result: 'mock2',
+          },
+        }),
+      );
+
+      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();
+    });
   });
 });