test: inherit from ERC20 to MockERC20

defi-wonderland · Jul 9, 2024 · 5d00a06 · 5d00a06
1 parent 2795bda
commit 5d00a06
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 16 deletions.
diff --git a/solidity/test/mocks/MockERC20.sol b/solidity/test/mocks/MockERC20.sol
@@ -1,15 +1,53 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.19;
 
-contract MockERC20 {
+import {ERC20} from '@openzeppelin/contracts/token/ERC20/ERC20.sol';
+
+contract MockERC20 is ERC20 {
   mapping(address _account => mapping(uint256 _callCount => uint256 _amount)) internal _balancesPerCall;
   mapping(address _account => uint256 _callCount) internal _callsPerAccount;
+  bool internal _mocked;
 
-  function balanceOf(address _account) external returns (uint256 _amount) {
-    _amount = _balancesPerCall[_account][_callsPerAccount[_account]++];
-  }
+  constructor(string memory _name, string memory _symbol) ERC20(_name, _symbol) {}
 
   function mockBalanceOfPerCall(address _account, uint256 _callCount, uint256 _amount) external {
     _balancesPerCall[_account][_callCount] = _amount;
+    _mocked = true;
+  }
+
+  function balanceOf(address _account) public view virtual override returns (uint256 _amount) {
+    if (_mocked) {
+      _amount = _balancesPerCall[_account][_callsPerAccount[_account]];
+    } else {
+      _amount = super.balanceOf(_account);
+    }
+  }
+
+  function transfer(address _to, uint256 _amount) public virtual override returns (bool _success) {
+    if (_mocked) {
+      ++_callsPerAccount[msg.sender];
+      ++_callsPerAccount[_to];
+    }
+    return super.transfer(_to, _amount);
+  }
+
+  function transferFrom(address _from, address _to, uint256 _amount) public virtual override returns (bool _success) {
+    if (_mocked) {
+      ++_callsPerAccount[_from];
+      ++_callsPerAccount[_to];
+    }
+    return super.transferFrom(_from, _to, _amount);
+  }
+
+  function mint(address _account, uint256 _amount) public virtual {
+    _mint(_account, _amount);
+  }
+
+  function burn(address _account, uint256 _amount) public virtual {
+    _burn(_account, _amount);
+  }
+
+  function approve(address _owner, address _spender, uint256 _amount) public virtual {
+    _approve(_owner, _spender, _amount);
   }
 }
diff --git a/solidity/test/unit/extensions/AccountingExtension.t.sol b/solidity/test/unit/extensions/AccountingExtension.t.sol
@@ -56,7 +56,7 @@ contract BaseTest is Test, Helpers {
     oracle = IOracle(makeAddr('Oracle'));
     vm.etch(address(oracle), hex'069420');
 
-    token = IERC20(address(new MockERC20()));
+    token = IERC20(new MockERC20('Token', 'TKN'));
 
     extension = new ForTest_AccountingExtension(oracle);
   }
@@ -67,16 +67,18 @@ contract AccountingExtension_Unit_DepositAndWithdraw is BaseTest {
    * @notice Test an ERC20 deposit
    */
   function test_depositERC20(uint256 _amount, uint256 _initialBalance) public {
+    // Mint tokens to the sender and approve the extension
+    MockERC20(address(token)).mint(sender, _amount);
+    MockERC20(address(token)).approve(sender, address(extension), _amount);
+
     // Mock and expect the ERC20 balance
     _initialBalance = bound(_initialBalance, 0, type(uint256).max - _amount);
     MockERC20(address(token)).mockBalanceOfPerCall(address(extension), 0, _initialBalance);
     MockERC20(address(token)).mockBalanceOfPerCall(address(extension), 1, _initialBalance + _amount);
     vm.expectCall(address(token), abi.encodeCall(IERC20.balanceOf, (address(extension))), 2);
 
-    // Mock and expect the ERC20 transfer
-    _mockAndExpect(
-      address(token), abi.encodeCall(IERC20.transferFrom, (sender, address(extension), _amount)), abi.encode(true)
-    );
+    // Expect the ERC20 transfer
+    vm.expectCall(address(token), abi.encodeCall(IERC20.transferFrom, (sender, address(extension), _amount)), 1);
 
     // Expect the event
     vm.expectEmit(true, true, true, true, address(extension));
@@ -96,16 +98,14 @@ contract AccountingExtension_Unit_DepositAndWithdraw is BaseTest {
     vm.assume(_amount >= _fee);
     vm.assume(_fee > 0);
 
-    // Mock and expect the ERC20 balance
+    // Mint tokens to the sender and approve the extension
+    MockERC20(address(token)).mint(sender, _amount);
+    MockERC20(address(token)).approve(sender, address(extension), _amount);
+
+    // Mock the ERC20 balance
     _initialBalance = bound(_initialBalance, 0, type(uint256).max - (_amount - _fee));
     MockERC20(address(token)).mockBalanceOfPerCall(address(extension), 0, _initialBalance);
     MockERC20(address(token)).mockBalanceOfPerCall(address(extension), 1, _initialBalance + (_amount - _fee));
-    vm.expectCall(address(token), abi.encodeCall(IERC20.balanceOf, (address(extension))), 2);
-
-    // Mock and expect the ERC20 transfer
-    _mockAndExpect(
-      address(token), abi.encodeCall(IERC20.transferFrom, (sender, address(extension), _amount)), abi.encode(true)
-    );
 
     // Check: does it revert if token takes a fee on transfer?
     vm.expectRevert(IAccountingExtension.AccountingExtension_FeeOnTransferToken.selector);
@@ -121,6 +121,7 @@ contract AccountingExtension_Unit_DepositAndWithdraw is BaseTest {
     _initialBalance = bound(_initialBalance, _amount, type(uint256).max);
     extension.forTest_setBalanceOf(sender, token, _initialBalance);
 
+    // Mock and expect the ERC20 transfer
     _mockAndExpect(address(token), abi.encodeCall(IERC20.transfer, (sender, _amount)), abi.encode(true));
 
     // Expect the event