From efac34bcb77d95b81af5bef44140895257b36231 Mon Sep 17 00:00:00 2001
From: Johannes Kares <jk@talentir.com>
Date: Tue, 6 Aug 2024 10:30:29 +0200
Subject: [PATCH] feat: Add `Wallet` to installMockWallet

Allow for a custom implementation of `Wallet` if default implementation is not sufficient
---
 src/installMockWallet.ts | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/src/installMockWallet.ts b/src/installMockWallet.ts
index d17abcb..3819396 100644
--- a/src/installMockWallet.ts
+++ b/src/installMockWallet.ts
@@ -6,26 +6,35 @@ import { randomUUID } from "crypto";
 let wallets: Map<string, Wallet> = new Map();
 
 export async function installMockWallet({
-  account,
-  transports,
-  defaultChain,
   debug,
   ...params
 }: {
-  account: LocalAccount;
-  transports?: Record<number, Transport>;
-  defaultChain?: Chain;
   debug?: boolean;
-} & ({ page: Page } | { browserContext: BrowserContext })) {
+} & ({ page: Page } | { browserContext: BrowserContext }) &
+  (
+    | {
+        account: LocalAccount;
+        transports?: Record<number, Transport>;
+        defaultChain?: Chain;
+      }
+    | {
+        wallet: Wallet;
+      }
+  )) {
   const browserOrPage =
     "browserContext" in params ? params.browserContext : params.page;
 
+  const wallet: Wallet =
+    "wallet" in params
+      ? params.wallet
+      : createWallet(params.account, params.transports, params.defaultChain);
+
   // Connecting the browser context to the Node.js playwright context
   await browserOrPage.exposeFunction("eip1193Request", eip1193Request);
 
   // Everytime we call installMockWallet, we create a new uuid to identify the wallet.
   const uuid = randomUUID();
-  wallets.set(uuid, createWallet(account, transports, defaultChain));
+  wallets.set(uuid, wallet);
 
   await browserOrPage.addInitScript(
     ({ uuid, debug }) => {