fix: use retryable contract to get token balance (#305)

# What ❔ Use retryable contract to get token balances. ## Why ❔ To use reties more efficient. ## Checklist   - [X] PR title corresponds to the body of PR (we generate changelog entries from PRs). - [X] Tests for the changes have been added / updated.
matter-labs · Oct 31, 2024 · 25606ec · 25606ec
1 parent 1cf7416
commit 25606ec
Show file tree

Hide file tree

Showing 4 changed files with 104 additions and 26 deletions.
diff --git a/packages/data-fetcher/src/blockchain/blockchain.service.spec.ts b/packages/data-fetcher/src/blockchain/blockchain.service.spec.ts
@@ -1884,7 +1884,7 @@ describe("BlockchainService", () => {
       it("gets the balance for ETH", async () => {
         await blockchainService.getBalance(address, blockNumber, tokenAddress);
         expect(provider.getBalance).toHaveBeenCalledTimes(1);
-        expect(provider.getBalance).toHaveBeenCalledWith(address, blockNumber, undefined);
+        expect(provider.getBalance).toHaveBeenCalledWith(address, blockNumber);
       });
 
       it("stops the rpc call duration metric", async () => {
@@ -2085,20 +2085,54 @@ describe("BlockchainService", () => {
     describe("if token address is not ETH", () => {
       beforeEach(() => {
         tokenAddress = "0x22b44df5aa1ee4542b6318ff971f183135f5e4ce";
-        jest.spyOn(provider, "getBalance").mockResolvedValue(BigInt(10));
       });
 
-      it("gets the balance for ETH", async () => {
-        await blockchainService.getBalance(address, blockNumber, tokenAddress);
-        expect(provider.getBalance).toHaveBeenCalledTimes(1);
-        expect(provider.getBalance).toHaveBeenCalledWith(address, blockNumber, tokenAddress);
+      describe("if ERC20 Contract function throws an exception", () => {
+        const error = new Error("Ethers Contract error");
+
+        beforeEach(() => {
+          (RetryableContract as any as jest.Mock).mockReturnValueOnce(
+            mock<RetryableContract>({
+              balanceOf: jest.fn().mockImplementationOnce(() => {
+                throw error;
+              }) as any,
+            })
+          );
+        });
+
+        it("throws an error", async () => {
+          await expect(blockchainService.getBalance(address, blockNumber, tokenAddress)).rejects.toThrowError(error);
+        });
       });
 
-      it("returns the address balance for ETH", async () => {
-        jest.spyOn(provider, "getBalance").mockResolvedValueOnce(BigInt(25));
+      describe("when there is a token with the specified address", () => {
+        let balanceOfMock: jest.Mock;
 
-        const balance = await blockchainService.getBalance(address, blockNumber, tokenAddress);
-        expect(balance).toStrictEqual(BigInt(25));
+        beforeEach(() => {
+          balanceOfMock = jest.fn().mockResolvedValueOnce(BigInt(20));
+          (RetryableContract as any as jest.Mock).mockReturnValueOnce(
+            mock<RetryableContract>({
+              balanceOf: balanceOfMock as any,
+            })
+          );
+        });
+
+        it("uses the proper token contract", async () => {
+          await blockchainService.getBalance(address, blockNumber, tokenAddress);
+          expect(RetryableContract).toHaveBeenCalledTimes(1);
+          expect(RetryableContract).toBeCalledWith(tokenAddress, utils.IERC20, provider);
+        });
+
+        it("gets the balance for the specified address and block", async () => {
+          await blockchainService.getBalance(address, blockNumber, tokenAddress);
+          expect(balanceOfMock).toHaveBeenCalledTimes(1);
+          expect(balanceOfMock).toHaveBeenCalledWith(address, { blockTag: blockNumber });
+        });
+
+        it("returns the balance of the token", async () => {
+          const balance = await blockchainService.getBalance(address, blockNumber, tokenAddress);
+          expect(balance).toStrictEqual(BigInt(20));
+        });
       });
     });
   });

diff --git a/packages/data-fetcher/src/blockchain/blockchain.service.ts b/packages/data-fetcher/src/blockchain/blockchain.service.ts
@@ -167,9 +167,14 @@ export class BlockchainService implements OnModuleInit {
   }
 
   public async getBalance(address: string, blockNumber: number, tokenAddress: string): Promise<bigint> {
-    return await this.rpcCall(async () => {
-      return await this.provider.getBalance(address, blockNumber, utils.isETH(tokenAddress) ? undefined : tokenAddress);
-    }, "getBalance");
+    if (utils.isETH(tokenAddress)) {
+      return await this.rpcCall(async () => {
+        return await this.provider.getBalance(address, blockNumber);
+      }, "getBalance");
+    }
+
+    const erc20Contract = new RetryableContract(tokenAddress, utils.IERC20, this.provider);
+    return await erc20Contract.balanceOf(address, { blockTag: blockNumber });
   }
 
   public async onModuleInit(): Promise<void> {

diff --git a/packages/worker/src/blockchain/blockchain.service.spec.ts b/packages/worker/src/blockchain/blockchain.service.spec.ts
@@ -1287,7 +1287,7 @@ describe("BlockchainService", () => {
       it("gets the balance for ETH", async () => {
         await blockchainService.getBalance(address, blockNumber, tokenAddress);
         expect(provider.getBalance).toHaveBeenCalledTimes(1);
-        expect(provider.getBalance).toHaveBeenCalledWith(address, blockNumber, undefined);
+        expect(provider.getBalance).toHaveBeenCalledWith(address, blockNumber);
       });
 
       it("stops the rpc call duration metric", async () => {
@@ -1413,20 +1413,54 @@ describe("BlockchainService", () => {
     describe("if token address is not ETH", () => {
       beforeEach(() => {
         tokenAddress = "0x22b44df5aa1ee4542b6318ff971f183135f5e4ce";
-        jest.spyOn(provider, "getBalance").mockResolvedValue(BigInt(10));
       });
 
-      it("gets the balance for ETH", async () => {
-        await blockchainService.getBalance(address, blockNumber, tokenAddress);
-        expect(provider.getBalance).toHaveBeenCalledTimes(1);
-        expect(provider.getBalance).toHaveBeenCalledWith(address, blockNumber, tokenAddress);
+      describe("if ERC20 Contract function throws an exception", () => {
+        const error = new Error("Ethers Contract error");
+
+        beforeEach(() => {
+          (RetryableContract as any as jest.Mock).mockReturnValueOnce(
+            mock<RetryableContract>({
+              balanceOf: jest.fn().mockImplementationOnce(() => {
+                throw error;
+              }) as any,
+            })
+          );
+        });
+
+        it("throws an error", async () => {
+          await expect(blockchainService.getBalance(address, blockNumber, tokenAddress)).rejects.toThrowError(error);
+        });
       });
 
-      it("returns the address balance for ETH", async () => {
-        jest.spyOn(provider, "getBalance").mockResolvedValueOnce(BigInt(25));
+      describe("when there is a token with the specified address", () => {
+        let balanceOfMock: jest.Mock;
 
-        const balance = await blockchainService.getBalance(address, blockNumber, tokenAddress);
-        expect(balance).toStrictEqual(BigInt(25));
+        beforeEach(() => {
+          balanceOfMock = jest.fn().mockResolvedValueOnce(BigInt(20));
+          (RetryableContract as any as jest.Mock).mockReturnValueOnce(
+            mock<RetryableContract>({
+              balanceOf: balanceOfMock as any,
+            })
+          );
+        });
+
+        it("uses the proper token contract", async () => {
+          await blockchainService.getBalance(address, blockNumber, tokenAddress);
+          expect(RetryableContract).toHaveBeenCalledTimes(1);
+          expect(RetryableContract).toBeCalledWith(tokenAddress, utils.IERC20, provider);
+        });
+
+        it("gets the balance for the specified address and block", async () => {
+          await blockchainService.getBalance(address, blockNumber, tokenAddress);
+          expect(balanceOfMock).toHaveBeenCalledTimes(1);
+          expect(balanceOfMock).toHaveBeenCalledWith(address, { blockTag: blockNumber });
+        });
+
+        it("returns the balance of the token", async () => {
+          const balance = await blockchainService.getBalance(address, blockNumber, tokenAddress);
+          expect(balance).toStrictEqual(BigInt(20));
+        });
       });
     });
   });

diff --git a/packages/worker/src/blockchain/blockchain.service.ts b/packages/worker/src/blockchain/blockchain.service.ts
@@ -159,9 +159,14 @@ export class BlockchainService implements OnModuleInit {
   }
 
   public async getBalance(address: string, blockNumber: number, tokenAddress: string): Promise<bigint> {
-    return await this.rpcCall(async () => {
-      return await this.provider.getBalance(address, blockNumber, utils.isETH(tokenAddress) ? undefined : tokenAddress);
-    }, "getBalance");
+    if (utils.isETH(tokenAddress)) {
+      return await this.rpcCall(async () => {
+        return await this.provider.getBalance(address, blockNumber);
+      }, "getBalance");
+    }
+
+    const erc20Contract = new RetryableContract(tokenAddress, utils.IERC20, this.provider);
+    return await erc20Contract.balanceOf(address, { blockTag: blockNumber });
   }
 
   public async onModuleInit(): Promise<void> {