BREAKING: Add support for new state methods to snaps-simulation

packages/examples/packages/manage-state/src/index.test.ts

@@ -20,6 +20,244 @@ describe('onRpcRequest', () => {
     });
   });
 
+  describe('setState', () => {
+    it('sets the state to the params', async () => {
+      const { request } = await installSnap();
+
+      expect(
+        await request({
+          method: 'setState',
+          params: {
+            value: {
+              items: ['foo'],
+            },
+          },
+        }),
+      ).toRespondWith(null);
+
+      expect(
+        await request({
+          method: 'getState',
+        }),
+      ).toRespondWith({
+        items: ['foo'],
+      });
+    });
+
+    it('sets the state at a specific key', async () => {
+      const { request } = await installSnap();
+
+      expect(
+        await request({
+          method: 'setState',
+          params: {
+            value: 'foo',
+            key: 'nested.key',
+          },
+        }),
+      ).toRespondWith(null);
+
+      expect(
+        await request({
+          method: 'getState',
+        }),
+      ).toRespondWith({
+        nested: {
+          key: 'foo',
+        },
+      });
+    });
+
+    it('sets the unencrypted state to the params', async () => {
+      const { request } = await installSnap();
+
+      expect(
+        await request({
+          method: 'setState',
+          params: {
+            value: {
+              items: ['foo'],
+            },
+            encrypted: false,
+          },
+        }),
+      ).toRespondWith(null);
+
+      expect(
+        await request({
+          method: 'getState',
+        }),
+      ).toRespondWith(null);
+
+      expect(
+        await request({
+          method: 'getState',
+          params: {
+            encrypted: false,
+          },
+        }),
+      ).toRespondWith({
+        items: ['foo'],
+      });
+    });
+
+    it('throws if the state is not an object and no key is specified', async () => {
+      const { request } = await installSnap();
+
+      const response = await request({
+        method: 'setState',
+        params: {
+          value: 'foo',
+        },
+      });
+
+      expect(response).toRespondWithError(
+        expect.objectContaining({
+          code: -32602,
+          message:
+            'Invalid params: Value must be an object if key is not provided.',
+        }),
+      );
+    });
+  });
+
+  describe('getState', () => {
+    it('returns `null` if no state has been set', async () => {
+      const { request } = await installSnap();
+
+      const response = await request({
+        method: 'getState',
+      });
+
+      expect(response).toRespondWith(null);
+    });
+
+    it('returns the state', async () => {
+      const { request } = await installSnap();
+
+      await request({
+        method: 'setState',
+        params: {
+          value: {
+            items: ['foo'],
+          },
+        },
+      });
+
+      const response = await request({
+        method: 'getState',
+      });
+
+      expect(response).toRespondWith({
+        items: ['foo'],
+      });
+    });
+
+    it('returns the state at a specific key', async () => {
+      const { request } = await installSnap();
+
+      await request({
+        method: 'setState',
+        params: {
+          value: {
+            nested: {
+              key: 'foo',
+            },
+          },
+        },
+      });
+
+      const response = await request({
+        method: 'getState',
+        params: {
+          key: 'nested.key',
+        },
+      });
+
+      expect(response).toRespondWith('foo');
+    });
+
+    it('returns the unencrypted state', async () => {
+      const { request } = await installSnap();
+
+      await request({
+        method: 'setState',
+        params: {
+          value: {
+            items: ['foo'],
+          },
+          encrypted: false,
+        },
+      });
+
+      const response = await request({
+        method: 'getState',
+        params: {
+          encrypted: false,
+        },
+      });
+
+      expect(response).toRespondWith({
+        items: ['foo'],
+      });
+    });
+  });
+
+  describe('clearState', () => {
+    it('clears the state', async () => {
+      const { request } = await installSnap();
+
+      await request({
+        method: 'setState',
+        params: {
+          value: {
+            items: ['foo'],
+          },
+        },
+      });
+
+      await request({
+        method: 'clearState',
+      });
+
+      const response = await request({
+        method: 'getState',
+      });
+
+      expect(response).toRespondWith(null);
+    });
+
+    it('clears the unencrypted state', async () => {
+      const { request } = await installSnap();
+
+      await request({
+        method: 'setState',
+        params: {
+          value: {
+            items: ['foo'],
+          },
+          encrypted: false,
+        },
+      });
+
+      await request({
+        method: 'clearState',
+        params: {
+          encrypted: false,
+        },
+      });
+
+      const response = await request({
+        method: 'getState',
+        params: {
+          encrypted: false,
+        },
+      });
+
+      expect(response).toRespondWith(null);
+    });
+  });
+
   describe('legacy_setState', () => {
     it('sets the state to the params', async () => {
       const { request } = await installSnap();

packages/snaps-simulation/src/controllers.test.ts

@@ -5,10 +5,10 @@ import {
 } from '@metamask/permission-controller';
 
 import { getControllers } from './controllers';
-import type { MiddlewareHooks } from './simulation';
+import type { RestrictedMiddlewareHooks } from './simulation';
 import { getMockOptions } from './test-utils';
 
-const MOCK_HOOKS: MiddlewareHooks = {
+const MOCK_HOOKS: RestrictedMiddlewareHooks = {
   getIsLocked: jest.fn(),
   getMnemonic: jest.fn(),
   getSnapFile: jest.fn(),

packages/snaps-simulation/src/controllers.ts

@@ -29,7 +29,7 @@ import { getSafeJson } from '@metamask/utils';
 import { getPermissionSpecifications } from './methods';
 import { UNRESTRICTED_METHODS } from './methods/constants';
 import type { SimulationOptions } from './options';
-import type { MiddlewareHooks } from './simulation';
+import type { RestrictedMiddlewareHooks } from './simulation';
 import type { RunSagaFunction } from './store';
 
 export type RootControllerAllowedActions =
@@ -49,7 +49,7 @@ export type RootControllerMessenger = ControllerMessenger<
 
 export type GetControllersOptions = {
   controllerMessenger: ControllerMessenger<any, any>;
-  hooks: MiddlewareHooks;
+  hooks: RestrictedMiddlewareHooks;
   runSaga: RunSagaFunction;
   options: SimulationOptions;
 };

packages/snaps-simulation/src/methods/hooks/index.ts

@@ -1,5 +1,6 @@
 export * from './get-preferences';
+export * from './interface';
 export * from './notifications';
+export * from './permitted';
 export * from './request-user-approval';
 export * from './state';
-export * from './interface';

packages/snaps-simulation/src/methods/hooks/permitted/index.ts

@@ -0,0 +1 @@
+export * from './state';