working tests

apps/extension/src/storage/base.ts

@@ -24,6 +24,7 @@ type Version = string;
 export type MigrationMap<OldState, NewState> = {
   [K in keyof OldState & keyof NewState]?: (
     prev: OldState[K],
+    get: <K extends keyof NewState>(key: K) => Promise<NewState[K]>,
   ) => NewState[K] | Promise<NewState[K]>;
 };
 
@@ -44,14 +45,28 @@ export class ExtensionStorage<T> {
   ) {}
 
   async get<K extends keyof T>(key: K): Promise<T[K]> {
+    return await this.getRaw({ key });
+  }
+
+  private async getRaw<K extends keyof T>({
+    key,
+    skipMigration = false,
+  }: {
+    key: K;
+    skipMigration?: boolean;
+  }): Promise<T[K]> {
     const result = (await this.storage.get(String(key))) as
       | Record<K, StorageItem<T[K]>>
       | EmptyObject;
 
     if (isEmptyObj(result)) {
       return this.defaults[key];
     } else {
-      return await this.migrateIfNeeded(key, result[key]);
+      if (skipMigration) {
+        return result[key].value;
+      } else {
+        return await this.migrateIfNeeded(key, result[key]);
+      }
     }
   }
 
@@ -85,7 +100,9 @@ export class ExtensionStorage<T> {
     const migrationFn = this.migrations[item.version]?.[key];
 
     // Perform migration if available for version & field
-    const value = migrationFn ? ((await migrationFn(item.value)) as T[K]) : item.value;
+    const value = migrationFn
+      ? await migrationFn(item.value, key => this.getRaw({ key, skipMigration: true }))
+      : item.value;
     const nextVersion = this.versionSteps[item.version];
 
     // If the next step is not defined (bad config) or is the current version, save and exit

apps/extension/src/storage/migrations/local-v2-migration.test.ts

@@ -5,6 +5,9 @@ import { localDefaults } from '../local';
 import { LocalStorageState, LocalStorageVersion } from '../types';
 import { localV1Migrations } from './local-v1-migrations';
 import { localV2Migrations } from './local-v2-migrations';
+import { ChainRegistryClient } from '@penumbra-labs/registry';
+import { sample } from 'lodash';
+import { AppParameters } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/app/v1/app_pb';
 
 describe('v2 local storage migrations', () => {
   const storageArea = new MockStorageArea();
@@ -26,6 +29,9 @@ describe('v2 local storage migrations', () => {
       {
         [LocalStorageVersion.V1]: localV1Migrations,
       },
+      {
+        [LocalStorageVersion.V1]: LocalStorageVersion.V2,
+      },
     );
 
     v3ExtStorage = new ExtensionStorage<LocalStorageState>(
@@ -36,6 +42,10 @@ describe('v2 local storage migrations', () => {
         [LocalStorageVersion.V1]: localV1Migrations,
         [LocalStorageVersion.V2]: localV2Migrations,
       },
+      {
+        [LocalStorageVersion.V1]: LocalStorageVersion.V2,
+        [LocalStorageVersion.V2]: LocalStorageVersion.V3,
+      },
     );
   });
 
@@ -47,11 +57,87 @@ describe('v2 local storage migrations', () => {
 
   describe('frontends', () => {
     test('not set frontend gets ignored', async () => {
-      await v1ExtStorage.set('frontendUrl', undefined);
+      await v1ExtStorage.set('frontendUrl', '');
+      const url = await v3ExtStorage.get('frontendUrl');
+      expect(url).toEqual('');
+    });
+
+    test('have no change if user already selected frontend in registry', async () => {
+      const registryClient = new ChainRegistryClient();
+      const { frontends } = registryClient.bundled.globals();
+      const suggestedFrontend = sample(frontends.map(f => f.url));
+      await v1ExtStorage.set('frontendUrl', suggestedFrontend);
+      const url = await v3ExtStorage.get('frontendUrl');
+      expect(url).toEqual(suggestedFrontend);
+    });
+
+    test('user gets migrated to suggested frontend', async () => {
+      const registryClient = new ChainRegistryClient();
+      const { frontends } = registryClient.bundled.globals();
+      await v1ExtStorage.set('frontendUrl', 'http://badfrontend.void');
+      const url = await v3ExtStorage.get('frontendUrl');
+      expect(url).not.toEqual('http://badfrontend.void');
+      expect(frontends.map(f => f.url).includes(url!)).toBeTruthy();
+    });
+
+    test('works from v2 storage as well', async () => {
+      const registryClient = new ChainRegistryClient();
+      const { frontends } = registryClient.bundled.globals();
+      await v2ExtStorage.set('frontendUrl', 'http://badfrontend.void');
       const url = await v3ExtStorage.get('frontendUrl');
-      expect(url).toBeUndefined();
+      expect(url).not.toEqual('http://badfrontend.void');
+      expect(frontends.map(f => f.url).includes(url!)).toBeTruthy();
     });
   });
 
-  test('Migration from v1 works the same', async () => {});
+  describe('grpcEndpoint', () => {
+    test('not set gets ignored', async () => {
+      await v1ExtStorage.set('grpcEndpoint', undefined);
+      const url = await v3ExtStorage.get('grpcEndpoint');
+      expect(url).toEqual(undefined);
+    });
+
+    test('not connected to mainnet gets ignored', async () => {
+      const appParams = new AppParameters({ chainId: 'testnet-deimos-42' });
+      await v1ExtStorage.set('params', appParams.toJsonString());
+      await v1ExtStorage.set('grpcEndpoint', 'grpc.testnet.void');
+      const endpoint = await v3ExtStorage.get('grpcEndpoint');
+      expect(endpoint).toEqual('grpc.testnet.void');
+    });
+
+    test('user selected suggested endpoint', async () => {
+      const appParams = new AppParameters({ chainId: 'penumbra-1' });
+      await v1ExtStorage.set('params', appParams.toJsonString());
+      const registryClient = new ChainRegistryClient();
+      const { rpcs } = registryClient.bundled.globals();
+      const suggestedRpc = sample(rpcs.map(f => f.url));
+      await v1ExtStorage.set('grpcEndpoint', suggestedRpc);
+      const endpoint = await v3ExtStorage.get('grpcEndpoint');
+      expect(endpoint).toEqual(suggestedRpc);
+    });
+
+    test('user gets migrated to suggested frontend', async () => {
+      const appParams = new AppParameters({ chainId: 'penumbra-1' });
+      await v1ExtStorage.set('params', appParams.toJsonString());
+      await v1ExtStorage.set('grpcEndpoint', 'http://badfrontend.void');
+      const endpoint = await v3ExtStorage.get('grpcEndpoint');
+      expect(endpoint).not.toEqual('http://badfrontend.void');
+
+      const registryClient = new ChainRegistryClient();
+      const { rpcs } = registryClient.bundled.globals();
+      expect(rpcs.map(r => r.url).includes(endpoint!)).toBeTruthy();
+    });
+
+    test('works from v2 storage as well', async () => {
+      const appParams = new AppParameters({ chainId: 'penumbra-1' });
+      await v2ExtStorage.set('params', appParams.toJsonString());
+      await v2ExtStorage.set('grpcEndpoint', 'http://badfrontend.void');
+      const endpoint = await v3ExtStorage.get('grpcEndpoint');
+      expect(endpoint).not.toEqual('http://badfrontend.void');
+
+      const registryClient = new ChainRegistryClient();
+      const { rpcs } = registryClient.bundled.globals();
+      expect(rpcs.map(r => r.url).includes(endpoint!)).toBeTruthy();
+    });
+  });
 });

apps/extension/src/storage/migrations/local-v2-migrations.ts

@@ -1,40 +1,35 @@
 import { LocalStorageState } from '../types';
 import { MigrationMap } from '../base';
-import { localExtStorage } from '../local';
 import { AppParameters } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/app/v1/app_pb';
 import { ChainRegistryClient } from '@penumbra-labs/registry';
 import { sample } from 'lodash';
 
 export const localV2Migrations: MigrationMap<LocalStorageState, LocalStorageState> = {
-  grpcEndpoint: async old => {
-    return await validateOrReplaceEndpoint(old);
+  grpcEndpoint: async (old, get) => {
+    return await validateOrReplaceEndpoint(old, get);
   },
   frontendUrl: old => {
     return validateOrReplaceFrontend(old);
   },
 };
 
-const isConnectedToMainnet = async (): Promise<boolean> => {
-  const chainId = await localExtStorage
-    .get('params')
-    .then(jsonParams =>
-      jsonParams ? AppParameters.fromJsonString(jsonParams).chainId : undefined,
-    );
-
-  // Ensure they are connected to mainnet
-  return chainId === 'penumbra-1';
-};
-
 // A one-time migration to suggested grpcUrls
 // Context: https://github.com/prax-wallet/web/issues/166
-const validateOrReplaceEndpoint = async (oldEndpoint?: string): Promise<string | undefined> => {
+const validateOrReplaceEndpoint = async (
+  oldEndpoint: string | undefined,
+  get: <K extends keyof LocalStorageState>(key: K) => Promise<LocalStorageState[K]>,
+): Promise<string | undefined> => {
   // If they don't have one set, it's likely they didn't go through onboarding
   if (!oldEndpoint) {
     return oldEndpoint;
   }
 
-  const connectedToMainnet = await isConnectedToMainnet();
-  if (!connectedToMainnet) {
+  // Ensure they are connected to mainnet
+  const chainId = await get('params').then(jsonParams =>
+    jsonParams ? AppParameters.fromJsonString(jsonParams).chainId : undefined,
+  );
+
+  if (chainId !== 'penumbra-1') {
     return oldEndpoint;
   }
 

apps/extension/src/storage/mock.ts

@@ -53,13 +53,11 @@ export const mockSessionExtStorage = () =>
     new MockStorageArea(),
     sessionDefaults,
     MockStorageVersion.V1,
-    {},
   );
 
 export const mockLocalExtStorage = () =>
   new ExtensionStorage<LocalStorageState>(
     new MockStorageArea(),
     localDefaults,
     MockStorageVersion.V1,
-    {},
   );

apps/extension/src/utils/tests-setup.js

@@ -7,7 +7,6 @@ global.chrome = {
     onChanged: {
       addListener: vi.fn(),
     },
-
     local: {
       set: vi.fn(),
       get: vi.fn().mockReturnValue({}),

docs/state-management.md

@@ -47,4 +47,4 @@ If your persisted state changes in a breaking way, it's important to write a mig
   }
 ```
 
-3. See [apps/extension/src/storage/migration.test.ts](../apps/extension/src/storage/base-migration.test.ts) for an example. Make sure you add types to your migration function!
+3. See [apps/extension/src/storage/migration.test.ts](../apps/extension/src/storage/migrations/base-migration.test.ts) for an example. Make sure you add types to your migration function!