From b620e545735a74bd9284608e029a75aee9b7f653 Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn <maarten@zuidhoorn.com>
Date: Tue, 10 Dec 2024 14:02:55 +0100
Subject: [PATCH] Add tests for new methods

---
 packages/snaps-rpc-methods/jest.config.js     |   8 +-
 .../src/permitted/clearState.test.ts          | 207 ++++++++++
 .../src/permitted/clearState.ts               |   3 +-
 .../src/permitted/getState.test.ts            | 281 +++++++++++++-
 .../src/permitted/getState.ts                 |   3 +-
 .../src/permitted/setState.test.ts            | 366 +++++++++++++++++-
 .../src/permitted/setState.ts                 |   4 +-
 7 files changed, 863 insertions(+), 9 deletions(-)
 create mode 100644 packages/snaps-rpc-methods/src/permitted/clearState.test.ts

diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js
index 077ed1c65d..30f84a3038 100644
--- a/packages/snaps-rpc-methods/jest.config.js
+++ b/packages/snaps-rpc-methods/jest.config.js
@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 92.94,
-      functions: 97.26,
-      lines: 97.87,
-      statements: 97.39,
+      branches: 93.3,
+      functions: 97.42,
+      lines: 98.01,
+      statements: 97.58,
     },
   },
 });
diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.test.ts b/packages/snaps-rpc-methods/src/permitted/clearState.test.ts
new file mode 100644
index 0000000000..1154341659
--- /dev/null
+++ b/packages/snaps-rpc-methods/src/permitted/clearState.test.ts
@@ -0,0 +1,207 @@
+import { JsonRpcEngine } from '@metamask/json-rpc-engine';
+import { errorCodes } from '@metamask/rpc-errors';
+import type { ClearStateResult } from '@metamask/snaps-sdk';
+import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils';
+import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils';
+
+import type { ClearStateParameters } from './clearState';
+import { clearStateHandler } from './clearState';
+
+describe('snap_clearState', () => {
+  describe('clearStateHandler', () => {
+    it('has the expected shape', () => {
+      expect(clearStateHandler).toMatchObject({
+        methodNames: ['snap_clearState'],
+        implementation: expect.any(Function),
+        hookNames: {
+          clearSnapState: true,
+          hasPermission: true,
+        },
+      });
+    });
+  });
+
+  describe('implementation', () => {
+    const createOriginMiddleware =
+      (origin: string) =>
+      (request: any, _response: unknown, next: () => void, _end: unknown) => {
+        request.origin = origin;
+        next();
+      };
+
+    it('returns the result from the `clearSnapState` hook', async () => {
+      const { implementation } = clearStateHandler;
+
+      const clearSnapState = jest.fn().mockReturnValue(null);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        clearSnapState,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<ClearStateParameters>,
+          response as PendingJsonRpcResponse<ClearStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_clearState',
+        params: {},
+      });
+
+      expect(clearSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true);
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+    });
+
+    it('clears unencrypted state if specified', async () => {
+      const { implementation } = clearStateHandler;
+
+      const clearSnapState = jest.fn().mockReturnValue(null);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        clearSnapState,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<ClearStateParameters>,
+          response as PendingJsonRpcResponse<ClearStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_clearState',
+        params: {
+          encrypted: false,
+        },
+      });
+
+      expect(clearSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, false);
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+    });
+
+    it('throws if the requesting origin does not have the required permission', async () => {
+      const { implementation } = clearStateHandler;
+
+      const clearSnapState = jest.fn();
+      const hasPermission = jest.fn().mockReturnValue(false);
+
+      const hooks = {
+        clearSnapState,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<ClearStateParameters>,
+          response as PendingJsonRpcResponse<ClearStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_clearState',
+        params: {},
+      });
+
+      expect(clearSnapState).not.toHaveBeenCalled();
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.provider.unauthorized,
+          message:
+            'The requested account and/or method has not been authorized by the user.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
+    it('throws if the parameters are invalid', async () => {
+      const { implementation } = clearStateHandler;
+
+      const clearSnapState = jest.fn();
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        clearSnapState,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<ClearStateParameters>,
+          response as PendingJsonRpcResponse<ClearStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_clearState',
+        params: {
+          encrypted: 'foo',
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: At path: encrypted -- Expected a value of type `boolean`, but received: `"foo"`.',
+          stack: expect.any(String),
+        },
+      });
+    });
+  });
+});
diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.ts b/packages/snaps-rpc-methods/src/permitted/clearState.ts
index bac8ff19f2..ae32a6ebe5 100644
--- a/packages/snaps-rpc-methods/src/permitted/clearState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/clearState.ts
@@ -117,6 +117,7 @@ function getValidatedParams(params?: unknown) {
       });
     }
 
-    throw error;
+    /* istanbul ignore next */
+    throw rpcErrors.internal();
   }
 }
diff --git a/packages/snaps-rpc-methods/src/permitted/getState.test.ts b/packages/snaps-rpc-methods/src/permitted/getState.test.ts
index 67a5df5a8d..252ab95d4d 100644
--- a/packages/snaps-rpc-methods/src/permitted/getState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/getState.test.ts
@@ -1,4 +1,283 @@
-import { get } from './getState';
+import { JsonRpcEngine } from '@metamask/json-rpc-engine';
+import { errorCodes } from '@metamask/rpc-errors';
+import type { GetStateResult } from '@metamask/snaps-sdk';
+import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils';
+import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils';
+
+import type { GetStateParameters } from './getState';
+import { get, getStateHandler } from './getState';
+
+describe('snap_getState', () => {
+  describe('getStateHandler', () => {
+    it('has the expected shape', () => {
+      expect(getStateHandler).toMatchObject({
+        methodNames: ['snap_getState'],
+        implementation: expect.any(Function),
+        hookNames: {
+          getSnapState: true,
+          hasPermission: true,
+        },
+      });
+    });
+  });
+
+  describe('implementation', () => {
+    const createOriginMiddleware =
+      (origin: string) =>
+      (request: any, _response: unknown, next: () => void, _end: unknown) => {
+        request.origin = origin;
+        next();
+      };
+
+    it('returns the encrypted state', async () => {
+      const { implementation } = getStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetStateParameters>,
+          response as PendingJsonRpcResponse<GetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getState',
+        params: {
+          key: 'foo',
+        },
+      });
+
+      expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true);
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: 'bar',
+      });
+    });
+
+    it('returns the entire state if no key is specified', async () => {
+      const { implementation } = getStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetStateParameters>,
+          response as PendingJsonRpcResponse<GetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getState',
+        params: {},
+      });
+
+      expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true);
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: {
+          foo: 'bar',
+        },
+      });
+    });
+
+    it('returns the unencrypted state', async () => {
+      const { implementation } = getStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetStateParameters>,
+          response as PendingJsonRpcResponse<GetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getState',
+        params: {
+          key: 'foo',
+          encrypted: false,
+        },
+      });
+
+      expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, false);
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: 'bar',
+      });
+    });
+
+    it('throws if the requesting origin does not have the required permission', async () => {
+      const { implementation } = getStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(false);
+
+      const hooks = {
+        getSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetStateParameters>,
+          response as PendingJsonRpcResponse<GetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getState',
+        params: {},
+      });
+
+      expect(getSnapState).not.toHaveBeenCalled();
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.provider.unauthorized,
+          message:
+            'The requested account and/or method has not been authorized by the user.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
+    it('throws if the parameters are invalid', async () => {
+      const { implementation } = getStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetStateParameters>,
+          response as PendingJsonRpcResponse<GetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getState',
+        params: {
+          encrypted: 'foo',
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: At path: encrypted -- Expected a value of type `boolean`, but received: `"foo"`.',
+          stack: expect.any(String),
+        },
+      });
+    });
+  });
+});
 
 describe('get', () => {
   const object = {
diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts
index 0018959e84..7ab1cd90f8 100644
--- a/packages/snaps-rpc-methods/src/permitted/getState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/getState.ts
@@ -137,7 +137,8 @@ function getValidatedParams(params?: unknown) {
       });
     }
 
-    throw error;
+    /* istanbul ignore next */
+    throw rpcErrors.internal();
   }
 }
 
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
index 2433e52015..5484bae711 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
@@ -1,4 +1,368 @@
-import { set } from './setState';
+import { JsonRpcEngine } from '@metamask/json-rpc-engine';
+import { errorCodes } from '@metamask/rpc-errors';
+import type { SetStateResult } from '@metamask/snaps-sdk';
+import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils';
+import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils';
+
+import { setStateHandler, type SetStateParameters, set } from './setState';
+
+describe('snap_setState', () => {
+  describe('setStateHandler', () => {
+    it('has the expected shape', () => {
+      expect(setStateHandler).toMatchObject({
+        methodNames: ['snap_setState'],
+        implementation: expect.any(Function),
+        hookNames: {
+          getSnapState: true,
+          hasPermission: true,
+        },
+      });
+    });
+  });
+
+  describe('implementation', () => {
+    const createOriginMiddleware =
+      (origin: string) =>
+      (request: any, _response: unknown, next: () => void, _end: unknown) => {
+        request.origin = origin;
+        next();
+      };
+
+    it('sets the encrypted state', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const updateSnapState = jest.fn().mockReturnValue(null);
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        updateSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          key: 'foo',
+          value: 'baz',
+        },
+      });
+
+      expect(getUnlockPromise).toHaveBeenCalled();
+      expect(updateSnapState).toHaveBeenCalledWith(
+        MOCK_SNAP_ID,
+        { foo: 'baz' },
+        true,
+      );
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+    });
+
+    it('sets the entire state if no key is specified', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const updateSnapState = jest.fn().mockReturnValue(null);
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        updateSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          value: {
+            foo: 'baz',
+          },
+        },
+      });
+
+      expect(getUnlockPromise).toHaveBeenCalled();
+      expect(updateSnapState).toHaveBeenCalledWith(
+        MOCK_SNAP_ID,
+        { foo: 'baz' },
+        true,
+      );
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+    });
+
+    it('sets the unencrypted state', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const updateSnapState = jest.fn().mockReturnValue(null);
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        updateSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          key: 'foo',
+          value: 'baz',
+          encrypted: false,
+        },
+      });
+
+      expect(getUnlockPromise).not.toHaveBeenCalled();
+      expect(updateSnapState).toHaveBeenCalledWith(
+        MOCK_SNAP_ID,
+        {
+          foo: 'baz',
+        },
+        false,
+      );
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+    });
+
+    it('throws if the requesting origin does not have the required permission', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const updateSnapState = jest.fn().mockReturnValue(null);
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(false);
+
+      const hooks = {
+        getSnapState,
+        updateSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {},
+      });
+
+      expect(updateSnapState).not.toHaveBeenCalled();
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.provider.unauthorized,
+          message:
+            'The requested account and/or method has not been authorized by the user.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
+    it('throws if the parameters are invalid', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const updateSnapState = jest.fn().mockReturnValue(null);
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        updateSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {},
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: At path: value -- Expected a value of type `JSON`, but received: `undefined`.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
+    it('throws if `key` is not provided and `value` is not an object', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const updateSnapState = jest.fn().mockReturnValue(null);
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        updateSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          value: 'foo',
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: Value must be an object if key is not provided.',
+          stack: expect.any(String),
+        },
+      });
+    });
+  });
+});
 
 describe('set', () => {
   it('sets the state in an empty object', () => {
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index f28463412d..64530320f2 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -171,7 +171,8 @@ function getValidatedParams(params?: unknown) {
       });
     }
 
-    throw error;
+    /* istanbul ignore next */
+    throw rpcErrors.internal();
   }
 }
 
@@ -225,5 +226,6 @@ export function set(
   }
 
   // This should never be reached.
+  /* istanbul ignore next */
   return {};
 }