feat(walletManager): add skipUsageRecording option and tests

Signed-off-by: Pablo Maldonado <pablo@umaproject.org>
UMAprotocol · Sep 18, 2024 · 5f05e0c · 5f05e0c
1 parent e607ef2
commit 5f05e0c
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 19 deletions.
diff --git a/src/lib/bundleUtils.ts b/src/lib/bundleUtils.ts
@@ -56,7 +56,7 @@ export const prepareUnlockTransaction = async (
 ) => {
   const provider = getProvider();
   const [baseFee, network] = await Promise.all([getBaseFee(provider, req), provider.getNetwork()]);
-  const unlockerWallet = WalletManager.getInstance().getWallet(ovalAddress, targetBlock, req.transactionId);
+  const unlockerWallet = WalletManager.getInstance().getWallet(ovalAddress, targetBlock, req.transactionId, simulate);
   const isSharedWallet = isOvalSharedUnlockerKey(unlockerWallet.address);
 
   // Encode the unlockLatestValue function call depending on whether the unlocker is a shared wallet or not.

diff --git a/src/lib/walletManager.ts b/src/lib/walletManager.ts
@@ -49,50 +49,56 @@ export class WalletManager {
     }
 
     // Get a wallet for a given address
-    public getWallet(address: string, targetBlock: number, transactionId: string): Wallet {
+    public getWallet(address: string, targetBlock: number, transactionId: string, skipUsageRecording = false): Wallet {
         if (!this.provider) {
             throw new Error("Provider is not initialized");
         }
         const checkSummedAddress = getAddress(address);
         const wallet = this.wallets[checkSummedAddress];
         if (!wallet) {
-            return this.getSharedWallet(address, targetBlock, transactionId);
+            return this.getSharedWallet(address, targetBlock, transactionId, skipUsageRecording);
         }
         return wallet.connect(this.provider);
     }
 
     // Get a shared wallet for a given Oval instance and target block
-    private getSharedWallet(ovalInstance: string, targetBlock: number, transactionId: string): Wallet {
+    private getSharedWallet(ovalInstance: string, targetBlock: number, transactionId: string, skipUsageRecording = false): Wallet {
         if (!this.provider) {
             throw new Error("Provider is not initialized");
         }
+
         if (!this.ovalDiscovery.isOval(ovalInstance)) {
             throw new Error(`Oval instance ${ovalInstance} is not found`);
         }
 
-        let selectedWallet: Wallet | undefined;
-
-        // Check if a wallet has already been assigned to this Oval instance
-        for (const [walletPubKey, instanceUsage] of this.sharedWalletUsage.entries()) {
-            for (const [_, record] of instanceUsage.entries()) {
-                if (record.ovalInstances && record.ovalInstances.has(ovalInstance)) {
-                    selectedWallet = this.sharedWallets.get(walletPubKey)!.connect(this.provider!);
-                }
-            }
-        }
+        let selectedWallet = this.findAssignedWallet(ovalInstance);
 
         // If no wallet has been assigned, find the least used wallet
         if (!selectedWallet) {
             selectedWallet = this.findLeastUsedWallet(transactionId);
         }
 
-        // Update the usage of the selected wallet
-        if (selectedWallet) {
+        if (!selectedWallet) {
+            throw new Error(`No available shared wallets for Oval instance ${ovalInstance} at block ${targetBlock}`);
+        }
+
+        // Update the usage of the selected wallet if recording is enabled
+        if (!skipUsageRecording) {
             this.updateWalletUsage(ovalInstance, selectedWallet, targetBlock);
-            return selectedWallet.connect(this.provider);
         }
 
-        throw new Error(`No available shared wallets for Oval instance ${ovalInstance} at block ${targetBlock}`);
+        return selectedWallet.connect(this.provider);
+    }
+
+    private findAssignedWallet(ovalInstance: string): Wallet | undefined {
+        for (const [walletPubKey, instanceUsage] of this.sharedWalletUsage.entries()) {
+            for (const record of instanceUsage.values()) {
+                if (record.ovalInstances?.has(ovalInstance)) {
+                    return this.sharedWallets.get(walletPubKey)?.connect(this.provider!);
+                }
+            }
+        }
+        return undefined;
     }
 
     public isOvalSharedUnlocker(unlockerPublicKey: string): boolean {
@@ -190,7 +196,6 @@ export class WalletManager {
         if (minInstances !== Infinity && minInstances !== 0) {
             Logger.error(transactionId, `Public key ${selectedWallet?.address} is reused in multiple Oval instances because no free wallets are available.`);
         }
-
         return selectedWallet;
     }
 

diff --git a/test/walletManager.ts b/test/walletManager.ts
@@ -141,6 +141,42 @@ describe('WalletManager Tests', () => {
         expect(errorSpy.args[0][1]).to.include(`Public key ${wallet2.address} is reused in multiple Oval instances because no free wallets are available.`);
     });
 
+    it('should not record usage when skipUsageRecording is true', async () => {
+        const unlockerRandom = getRandomAddressAndKey();
+        const sharedConfigs: OvalConfigsShared = [
+            { unlockerKey: unlockerRandom.privateKey },
+        ];
+        const walletManager = WalletManager.getInstance();
+        await walletManager.initialize(mockProvider, {}, sharedConfigs);
+
+        const ovalInstance = getRandomAddressAndKey().address;
+        const targetBlock = 123;
+
+        const wallet = walletManager.getWallet(ovalInstance, targetBlock, "transactionId", true);
+
+        const sharedWalletUsage = walletManager['sharedWalletUsage'].get(await wallet.getAddress());
+        expect(sharedWalletUsage).to.be.undefined;
+    });
+
+    it('should record usage when skipUsageRecording is false', async () => {
+        const unlockerRandom = getRandomAddressAndKey();
+        const sharedConfigs: OvalConfigsShared = [
+            { unlockerKey: unlockerRandom.privateKey },
+        ];
+        const walletManager = WalletManager.getInstance();
+        await walletManager.initialize(mockProvider, {}, sharedConfigs);
+
+        const ovalInstance = getRandomAddressAndKey().address;
+        const targetBlock = 123;
+
+        const wallet = await walletManager.getWallet(ovalInstance, targetBlock, "transactionId");
+
+        const sharedWalletUsage = walletManager['sharedWalletUsage'].get(await wallet.getAddress());
+        expect(sharedWalletUsage).to.not.be.undefined;
+        expect(sharedWalletUsage?.get(targetBlock)?.count).to.equal(1);
+        expect(sharedWalletUsage?.get(targetBlock)?.ovalInstances.has(ovalInstance)).to.be.true;
+    });
+
     it('should cleanup old records correctly', async () => {
         const unlockerRandom = getRandomAddressAndKey();
         const sharedConfigs: OvalConfigsShared = [