interledger · sidvishnoi · Jan 8, 2025 · Jan 7, 2025
diff --git a/src/shared/types.ts b/src/shared/types.ts
@@ -144,6 +144,10 @@ export type DeepReadonly<T> = {
   readonly [P in keyof T]: DeepReadonly<T[P]>;
 };
 
+export type DeepPartial<T> = T extends object
+  ? { [P in keyof T]?: DeepPartial<T[P]> }
+  : T;
+
 export type RequiredFields<T, K extends keyof T> = T & Required<Pick<T, K>>;
 
 export type Tab = RequiredFields<Tabs.Tab, 'id' | 'url'>;

diff --git a/tests/e2e/basic.spec.ts b/tests/e2e/basic.spec.ts
@@ -41,11 +41,7 @@ test.describe('should fail to connect if:', () => {
       i18n.getMessage('connectWallet_error_urlInvalidUrl'),
     );
 
-    expect(
-      await background.evaluate(() => {
-        return chrome.storage.local.get(['connected']);
-      }),
-    ).toEqual({ connected: false });
+    await expect(background).toHaveStorage({ connected: false });
   });
 
   test('invalid wallet address provided', async ({
@@ -62,10 +58,6 @@ test.describe('should fail to connect if:', () => {
       'not a valid wallet address',
     );
 
-    expect(
-      await background.evaluate(() => {
-        return chrome.storage.local.get(['connected']);
-      }),
-    ).toEqual({ connected: false });
+    await expect(background).toHaveStorage({ connected: false });
   });
 });
diff --git a/tests/e2e/connect.spec.ts b/tests/e2e/connect.spec.ts
@@ -24,11 +24,7 @@ test('connects with correct details provided', async ({
   expect(TEST_WALLET_PUBLIC_KEY).toBeDefined();
   expect(TEST_WALLET_ADDRESS_URL).toBeDefined();
 
-  expect(
-    await background.evaluate(() => {
-      return chrome.storage.local.get(['connected']);
-    }),
-  ).toEqual({ connected: false });
+  await expect(background).toHaveStorage({ connected: false });
 
   const keyInfo = {
     keyId: TEST_WALLET_KEY_ID,
@@ -45,14 +41,7 @@ test('connects with correct details provided', async ({
   const settingsLink = popup.locator(`[href="/settings"]`).first();
   await expect(settingsLink).toBeVisible();
 
-  const storage = await background.evaluate(() => {
-    return chrome.storage.local.get([
-      'connected',
-      'oneTimeGrant',
-      'recurringGrant',
-    ]);
-  });
-  expect(storage).toMatchObject({
+  await expect(background).toHaveStorage({
     connected: true,
     recurringGrant: null,
     oneTimeGrant: {
@@ -69,9 +58,5 @@ test('connects with correct details provided', async ({
   );
 
   await disconnectWallet(popup);
-  expect(
-    await background.evaluate(() => {
-      return chrome.storage.local.get(['connected']);
-    }),
-  ).toEqual({ connected: false });
+  await expect(background).toHaveStorage({ connected: false });
 });
diff --git a/tests/e2e/connectAutoKeyChimoney.spec.ts b/tests/e2e/connectAutoKeyChimoney.spec.ts
@@ -7,6 +7,7 @@ import {
   revokeKey,
   waitForGrantConsentPage,
 } from './helpers/chimoney';
+import { getStorage } from './fixtures/helpers';
 
 test('Connect to Chimoney wallet with automatic key addition when not logged-in to wallet', async ({
   page,
@@ -28,9 +29,7 @@ test('Connect to Chimoney wallet with automatic key addition when not logged-in
   test.slow(true, 'Some pages load slow');
 
   const walletURL = new URL(walletUrl);
-  const { keyId } = await background.evaluate(() => {
-    return chrome.storage.local.get<{ keyId: string }>(['keyId']);
-  });
+  const { keyId } = await getStorage(background, ['keyId']);
 
   const connectButton = await test.step('fill popup', async () => {
     const connectButton = await fillPopup(popup, i18n, {

diff --git a/tests/e2e/connectAutoKeyFynbos.spec.ts b/tests/e2e/connectAutoKeyFynbos.spec.ts
@@ -9,6 +9,7 @@ import {
   revokeKey,
   waitForGrantConsentPage,
 } from './helpers/fynbos';
+import { getStorage } from './fixtures/helpers';
 
 test('Connect to Fynbos with automatic key addition when not logged-in to wallet', async ({
   page,
@@ -23,9 +24,7 @@ test('Connect to Fynbos with automatic key addition when not logged-in to wallet
 
   test.skip(!username || !password || !walletAddressUrl, 'Missing credentials');
 
-  const { keyId: kid } = await background.evaluate(() => {
-    return chrome.storage.local.get<{ keyId: string }>(['keyId']);
-  });
+  const { keyId: kid } = await getStorage(background, ['keyId']);
 
   const connectButton = await test.step('fill popup', async () => {
     const connectButton = await fillPopup(popup, i18n, {
@@ -119,9 +118,7 @@ test('Connect to Fynbos with automatic key addition when not logged-in to wallet
     await acceptGrant(page, continueWaitMs);
     await waitForWelcomePage(page);
 
-    expect(
-      await background.evaluate(() => chrome.storage.local.get(['connected'])),
-    ).toEqual({ connected: true });
+    expect(background).toHaveStorage({ connected: true });
   });
 
   await test.step('cleanup: revoke keys', async () => {

diff --git a/tests/e2e/connectAutoKeyTestWallet.spec.ts b/tests/e2e/connectAutoKeyTestWallet.spec.ts
@@ -11,6 +11,7 @@ import {
   revokeKey,
   waitForGrantConsentPage,
 } from './helpers/testWallet';
+import { getStorage } from './fixtures/helpers';
 
 test('Connect to test wallet with automatic key addition when not logged-in to wallet', async ({
   page,
@@ -106,9 +107,7 @@ test('Connect to test wallet with automatic key addition when not logged-in to w
       }
     });
 
-    const { keyId } = await background.evaluate(() => {
-      return chrome.storage.local.get<{ keyId: string }>(['keyId']);
-    });
+    const { keyId } = await getStorage(background, ['keyId']);
 
     const jwksBefore = await getJWKS(walletAddressUrl);
     expect(jwksBefore.keys.length).toBeGreaterThanOrEqual(0);
@@ -137,9 +136,7 @@ test('Connect to test wallet with automatic key addition when not logged-in to w
     await acceptGrant(page, continueWaitMs);
     await waitForWelcomePage(page);
 
-    expect(
-      await background.evaluate(() => chrome.storage.local.get(['connected'])),
-    ).toEqual({ connected: true });
+    await expect(background).toHaveStorage({ connected: true });
   });
 
   await test.step('revoke key', async () => {

diff --git a/tests/e2e/fixtures/base.ts b/tests/e2e/fixtures/base.ts
@@ -1,11 +1,13 @@
 import { test as base, type BrowserContext, type Page } from '@playwright/test';
 import {
   getBackground,
+  getStorage,
   loadContext,
   BrowserIntl,
   type Background,
 } from './helpers';
 import { openPopup, type Popup } from '../pages/popup';
+import type { DeepPartial, Storage } from '@/shared/types';
 
 type BaseScopeWorker = {
   persistentContext: BrowserContext;
@@ -66,4 +68,51 @@ export const test = base.extend<{ page: Page }, BaseScopeWorker>({
   },
 });
 
-export const expect = test.expect;
+export const expect = test.expect.extend({
+  async toHaveStorage(background: Background, expected: DeepPartial<Storage>) {
+    const assertionName = 'toHaveStorage';
+
+    let pass: boolean;
+    let matcherResult: any;
+
+    const storedData = await getStorage(
+      background,
+      Object.keys(expected) as Array<keyof typeof expected>,
+    );
+    try {
+      test.expect(storedData).toMatchObject(expected);
+      pass = true;
+    } catch {
+      matcherResult = { actual: storedData };
+      pass = false;
+    }
+
+    const message = pass
+      ? () =>
+          this.utils.matcherHint(assertionName, undefined, undefined, {
+            isNot: this.isNot,
+          }) +
+          '\n\n' +
+          `Expected: not ${this.utils.printExpected(expected)}\n` +
+          (matcherResult
+            ? `Received: ${this.utils.printReceived(matcherResult.actual)}`
+            : '')
+      : () =>
+          this.utils.matcherHint(assertionName, undefined, undefined, {
+            isNot: this.isNot,
+          }) +
+          '\n\n' +
+          `Expected: ${this.utils.printExpected(expected)}\n` +
+          (matcherResult
+            ? `Received: ${this.utils.printReceived(matcherResult.actual)}`
+            : '');
+
+    return {
+      name: assertionName,
+      pass,
+      expected,
+      actual: matcherResult?.actual,
+      message,
+    };
+  },
+});
diff --git a/tests/e2e/fixtures/connected.ts b/tests/e2e/fixtures/connected.ts
@@ -1,5 +1,5 @@
-import type { Page } from '@playwright/test';
-import { test as base } from './base';
+import { mergeExpects, type Page } from '@playwright/test';
+import { test as base, expect as baseExpect } from './base';
 import { connectWallet, disconnectWallet, type Popup } from '../pages/popup';
 
 // With extension connected to the wallet.
@@ -26,4 +26,4 @@ export const test = base.extend<{ page: Page }, { popup: Popup }>({
   ],
 });
 
-export const expect = test.expect;
+export const expect = mergeExpects(baseExpect, test.expect);
diff --git a/tests/e2e/fixtures/helpers.ts b/tests/e2e/fixtures/helpers.ts
@@ -15,6 +15,7 @@ import {
 import { APP_URL } from '@/background/constants';
 import { DIST_DIR, ROOT_DIR } from '../../../esbuild/config';
 import type { TranslationKeys } from '../../../src/shared/helpers';
+import type { Storage, StorageKey } from '../../../src/shared/types';
 
 export type BrowserInfo = { browserName: string; channel: string | undefined };
 export type Background = Worker;
@@ -281,9 +282,11 @@ export async function loadKeysToExtension(
     });
   }, keyInfo);
 
-  const res = await background.evaluate(() => {
-    return chrome.storage.local.get(['privateKey', 'publicKey', 'keyId']);
-  });
+  const res = await getStorage(background, [
+    'privateKey',
+    'publicKey',
+    'keyId',
+  ]);
   if (!res || !res.keyId || !res.privateKey || !res.publicKey) {
     throw new Error('Could not load keys to extension');
   }
@@ -345,3 +348,14 @@ export class BrowserIntl {
     return result;
   }
 }
+
+export async function getStorage<TKey extends StorageKey>(
+  background: Background,
+  keys?: TKey[],
+) {
+  const data = await background.evaluate(
+    (keys) => chrome.storage.local.get(keys),
+    keys,
+  );
+  return data as { [Key in TKey[][number]]: Storage[Key] };
+}