From 66c7e254efc3a67cfa793c33f7bff3d69acf0e22 Mon Sep 17 00:00:00 2001 From: Vasyl Ivanchuk Date: Thu, 31 Oct 2024 21:02:44 +0200 Subject: [PATCH] fix: use retryable contract to get token balance --- .../src/blockchain/blockchain.service.spec.ts | 54 +++++++++++++++---- .../src/blockchain/blockchain.service.ts | 11 ++-- .../src/blockchain/blockchain.service.spec.ts | 54 +++++++++++++++---- .../src/blockchain/blockchain.service.ts | 11 ++-- 4 files changed, 104 insertions(+), 26 deletions(-) diff --git a/packages/data-fetcher/src/blockchain/blockchain.service.spec.ts b/packages/data-fetcher/src/blockchain/blockchain.service.spec.ts index 3cdbf817b..bd0f65d25 100644 --- 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({ + 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({ + 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 index 7ee73da70..6e4cd86f6 100644 --- 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 { - 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 { diff --git a/packages/worker/src/blockchain/blockchain.service.spec.ts b/packages/worker/src/blockchain/blockchain.service.spec.ts index eb636decb..e93bd1c1e 100644 --- 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({ + 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({ + 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 index 2303614ba..3bef4cd8f 100644 --- 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 { - 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 {