Add gasUsed test + refactor

zodenode · May 2, 2023 · 176eee3 · 176eee3
1 parent 0bb3bbb
commit 176eee3
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 12 deletions.
diff --git a/contracts/tests/mocks/SafeERC20Helper.sol b/contracts/tests/mocks/SafeERC20Helper.sol
@@ -13,6 +13,11 @@ contract ERC20ReturnFalseMock {
     // we write to a dummy state variable.
     uint256 private _dummy;
 
+    function balanceOf(address) external view returns (uint256) {
+        require(_dummy == 0, "Dummy"); // Dummy read from a state variable so that the function is view
+        return 0;
+    }
+
     function transfer(address, uint256) public returns (bool) {
         _dummy = 0;
         return false;
@@ -178,6 +183,10 @@ contract SafeERC20Wrapper {
         _token = token;
     }
 
+    function balanceOf() public view returns (uint256) {
+        return _token.safeBalanceOf(address(0));
+    }
+
     function transfer() public {
         _token.safeTransfer(address(0), 0);
     }
@@ -248,3 +257,21 @@ contract SafeWETHWrapper {
         _token.safeWithdrawTo(amount, to);
     }
 }
+
+contract ERC20WithSafeBalance {
+    using SafeERC20 for IERC20;
+
+    IERC20 private immutable _token;
+
+    constructor(IERC20 token) {
+        _token = token;
+    }
+
+    function balanceOf(address account) public view returns (uint256) {
+        return _token.balanceOf(account);
+    }
+
+    function safeBalanceOf(address account) public view returns (uint256) {
+        return _token.safeBalanceOf(account);
+    }
+}
diff --git a/test/contracts/SafestERC20.test.ts b/test/contracts/SafestERC20.test.ts
@@ -120,16 +120,36 @@ describe('SafeERC20', function () {
         return { weth, wrapper };
     }
 
+    async function deployERC20WithSafeBalance() {
+        const WETH = await ethers.getContractFactory('WETH');
+        const weth = await WETH.deploy();
+        await weth.deployed();
+
+        const ERC20WithSafeBalance = await ethers.getContractFactory('ERC20WithSafeBalance');
+        const wrapper = await ERC20WithSafeBalance.deploy(weth.address);
+        await wrapper.deployed();
+        return { weth, wrapper };
+    }
+
     describe('with address that has no contract code', function () {
         shouldRevertOnAllCalls(
-            ['SafeTransferFailed', 'SafeTransferFromFailed', 'ForceApproveFailed', ''],
+            {
+                transfer: 'SafeTransferFailed',
+                transferFrom: 'SafeTransferFromFailed',
+                approve: 'ForceApproveFailed',
+                changeAllowance: '',
+            },
             deployWrapperSimple,
         );
     });
 
     describe('with token that returns false on all calls', function () {
         shouldRevertOnAllCalls(
-            ['SafeTransferFailed', 'SafeTransferFromFailed', 'ForceApproveFailed'],
+            {
+                transfer: 'SafeTransferFailed',
+                transferFrom: 'SafeTransferFromFailed',
+                approve: 'ForceApproveFailed',
+            },
             deployWrapperFalseMock,
         );
     });
@@ -162,6 +182,22 @@ describe('SafeERC20', function () {
         });
     });
 
+    describe('safeBalanceOf', function () {
+        it('should be cheaper than balanceOf', async function () {
+            const { wrapper } = await loadFixture(deployERC20WithSafeBalance);
+
+            const tx = await wrapper.populateTransaction.balanceOf(owner.address);
+            const receipt = await owner.sendTransaction(tx);
+            const gasUsed = (await receipt.wait()).gasUsed;
+            const safeTx = await wrapper.populateTransaction.safeBalanceOf(owner.address);
+            const safeReceipt = await owner.sendTransaction(safeTx);
+            const safeGasUsed = (await safeReceipt.wait()).gasUsed;
+
+            expect(gasUsed).gt(safeGasUsed);
+            console.log(`balanceOf:safeBalanceOf gasUsed - ${gasUsed.toString()}:${safeGasUsed.toString()}`);
+        });
+    });
+
     describe("with token that doesn't revert on invalid permit", function () {
         it('accepts owner signature', async function () {
             const { token, wrapper, data, signature } = await loadFixture(deployPermitNoRevertAndSign);
@@ -311,37 +347,37 @@ describe('SafeERC20', function () {
         });
     });
 
-    function shouldRevertOnAllCalls(reasons: string[], fixture: () => Promise<{ wrapper: Contract }>) {
+    function shouldRevertOnAllCalls(reasons: { [methodName: string]: string }, fixture: () => Promise<{ wrapper: Contract }>) {
         it('reverts on transfer', async function () {
             const { wrapper } = await loadFixture(fixture);
-            await expect(wrapper.transfer()).to.be.revertedWithCustomError(wrapper, reasons[0]);
+            await expect(wrapper.transfer()).to.be.revertedWithCustomError(wrapper, reasons.transfer);
         });
 
         it('reverts on transferFrom', async function () {
             const { wrapper } = await loadFixture(fixture);
-            await expect(wrapper.transferFrom()).to.be.revertedWithCustomError(wrapper, reasons[1]);
+            await expect(wrapper.transferFrom()).to.be.revertedWithCustomError(wrapper, reasons.transferFrom);
         });
 
         it('reverts on approve', async function () {
             const { wrapper } = await loadFixture(fixture);
-            await expect(wrapper.approve(0)).to.be.revertedWithCustomError(wrapper, reasons[2]);
+            await expect(wrapper.approve(0)).to.be.revertedWithCustomError(wrapper, reasons.approve);
         });
 
         it('reverts on increaseAllowance', async function () {
             const { wrapper } = await loadFixture(fixture);
-            if (reasons.length === 3) {
-                await expect(wrapper.increaseAllowance(0)).to.be.revertedWithCustomError(wrapper, reasons[2]);
-            } else {
+            if (reasons.changeAllowance === '') {
                 await expect(wrapper.increaseAllowance(0)).to.be.reverted;
+            } else {
+                await expect(wrapper.increaseAllowance(0)).to.be.revertedWithCustomError(wrapper, reasons.approve);
             }
         });
 
         it('reverts on decreaseAllowance', async function () {
             const { wrapper } = await loadFixture(fixture);
-            if (reasons.length === 3) {
-                await expect(wrapper.decreaseAllowance(0)).to.be.revertedWithCustomError(wrapper, reasons[2]);
-            } else {
+            if (reasons.changeAllowance === '') {
                 await expect(wrapper.decreaseAllowance(0)).to.be.reverted;
+            } else {
+                await expect(wrapper.decreaseAllowance(0)).to.be.revertedWithCustomError(wrapper, reasons.approve);
             }
         });
     }