From 8e3a59c1753d4de5ed209f8fa163a369c5b86278 Mon Sep 17 00:00:00 2001
From: solanoe <solanoepalacio@gmail.com>
Date: Tue, 5 Dec 2023 23:26:08 -0300
Subject: [PATCH] refactor pullBalancesAtBlockHeight

---
 src/wallet-manager.ts | 32 ++++++++++----------------------
 1 file changed, 10 insertions(+), 22 deletions(-)

diff --git a/src/wallet-manager.ts b/src/wallet-manager.ts
index 27e6845..6b2ff52 100644
--- a/src/wallet-manager.ts
+++ b/src/wallet-manager.ts
@@ -429,29 +429,17 @@ export class WalletManager {
     }, Object.keys(this.managers).length);
   }
 
-  // pullBalancesAtBlockHeight doesn't need balances to be refreshed in the background
-  public async pullBalancesAtBlockHeight(
-    blockHeightByChain?: Record<ChainName, number>,
-  ): Promise<Record<string, WalletBalancesByAddress>> {
-    const balances: Record<string, WalletBalancesByAddress> = {};
-    if (blockHeightByChain) {
-      this.validateBlockHeightByChain(blockHeightByChain);
-    }
-
-    const blockHeightPerChain = blockHeightByChain ?? await this.getBlockHeightForAllSupportedChains();
-
-    await mapConcurrent(
-      Object.entries(this.managers),
-      async ([chainName, manager]) => {
-        const blockHeight = blockHeightPerChain[chainName as ChainName];
-        const balancesByChain = await manager.pullBalancesAtBlockHeight(
-          blockHeight,
-        );
-        balances[chainName] = balancesByChain;
-      },
-    );
+  // This method is similar to pullBalances, but it gets the block-height of each chain concurrently 
+  // and uses the resulting block-height to pull balances at that specific block-height
+  // it tends to return a balance that better represents a snapshot in time (as opposed to pullBalances
+  // which tries to return the latest possible balance)
+  public async pullBalancesAtCurrentBlockHeight(): Promise<Record<string, WalletBalancesByAddress>> {
+    const blockHeightByChain = await this.getBlockHeightForAllSupportedChains();
 
-    return balances;
+    return this.mapToChains(async (chainName: ChainName, manager: ChainWalletManager) => {
+      const blockHeight = blockHeightByChain[chainName];
+      return manager.pullBalancesAtBlockHeight(blockHeight);  
+    });
   }
 
   public getChainBalances(chainName: ChainName): WalletBalancesByAddress {