From b8bdb4dc85d44493a239dbadea6b26519bbab519 Mon Sep 17 00:00:00 2001 From: Dzmitry Lahunouski Date: Mon, 22 Jul 2024 21:36:35 +0100 Subject: [PATCH] unit tests --- src/Miller.sol | 2 +- test/Miller.t.sol | 106 ++++++++++++++++++++++++++++++++++++--- test/mocks/MockERC20.sol | 12 +++++ test/utils/SigUtils.sol | 43 ++++++++++++++++ 4 files changed, 155 insertions(+), 8 deletions(-) create mode 100644 test/mocks/MockERC20.sol create mode 100644 test/utils/SigUtils.sol diff --git a/src/Miller.sol b/src/Miller.sol index 7f3fe69..ff87a95 100644 --- a/src/Miller.sol +++ b/src/Miller.sol @@ -109,7 +109,7 @@ contract Miller is Context { } function _withdrawERC20(address to, IERC20 erc20token, uint240 amount) private { - erc20token.safeTransfer(to, amount); + erc20token.safeTransferFrom(_msgSender(), to, amount); } function _safePermit( diff --git a/test/Miller.t.sol b/test/Miller.t.sol index e36cd9d..a5045ec 100644 --- a/test/Miller.t.sol +++ b/test/Miller.t.sol @@ -10,17 +10,41 @@ import "@std/Test.sol"; import {Strings} from "@openzeppelin/contracts/utils/Strings.sol"; import {Miller} from "src/Miller.sol"; +import {MockERC20} from "./mocks/MockERC20.sol"; +import {SigUtils} from "./utils/SigUtils.sol"; + contract MillerTest is Test { Miller private miller; - address private alice = makeAddr("alice"); + address private Alice; + uint256 private AlicePk; + MockERC20 private mockERC20 = new MockERC20(); + SigUtils internal sigUtils = new SigUtils(mockERC20.DOMAIN_SEPARATOR()); function setUp() public { + (Alice, AlicePk) = makeAddrAndKey("alice"); miller = new Miller(); - vm.deal(alice, 10 ether); + vm.deal(Alice, 10 ether); + } + + function _simplePermit(address account, uint256 amount) + internal + view + returns (uint8, bytes32, bytes32) + { + SigUtils.Permit memory permit = SigUtils.Permit({ + owner: account, + spender: address(miller), + value: amount, + nonce: 0, + deadline: 1 days + }); + + bytes32 digest = sigUtils.getTypedDataHash(permit); + return vm.sign(AlicePk, digest); } - function generateAccounts(uint8 amount) private returns (address[] memory) { + function _generateAccounts(uint8 amount) private returns (address[] memory) { address[] memory addressList = new address[](amount); for (uint256 i = 0; i < addressList.length; i++) { addressList[i] = makeAddr(Strings.toString(i * 20)); @@ -28,16 +52,84 @@ contract MillerTest is Test { return addressList; } + function _generateConfig(uint8 amount) + private + returns (Miller.DistributionConfig[] memory, uint240) + { + Miller.DistributionConfig[] memory config = new Miller.DistributionConfig[](amount); + uint240 totalToDistribute; + for (uint256 i = 0; i < config.length; i++) { + uint240 amountToDistribute = uint240(i) * 20; + totalToDistribute += amountToDistribute; + config[i] = Miller.DistributionConfig({ + to: makeAddr(Strings.toString(i * 20)), + amount: amountToDistribute + }); + } + return (config, totalToDistribute); + } + function testFuzz_distributeFixed(uint8 addressesAmount, uint32 amountToDistribute) public { vm.assume(addressesAmount > 0); uint240 totalDistribute = uint240(amountToDistribute) * uint240(addressesAmount); - address[] memory addressList = generateAccounts(addressesAmount); - console.log("address %s", addressList[0]); - vm.prank(alice); + address[] memory addressList = _generateAccounts(addressesAmount); + vm.prank(Alice); miller.distributeFixed{value: totalDistribute}(uint240(amountToDistribute), addressList); - assertEq(alice.balance, 10 ether - uint256(totalDistribute)); + assertEq(Alice.balance, 10 ether - uint256(totalDistribute)); for (uint256 i = 0; i < addressList.length; i++) { assertEq(addressList[i].balance, amountToDistribute); } } + + function testFuzz_distribute(uint8 addressesAmount) public { + vm.assume(addressesAmount > 0); + (Miller.DistributionConfig[] memory config, uint240 totalDistribute) = + _generateConfig(addressesAmount); + vm.prank(Alice); + miller.distribute{value: totalDistribute}(config); + assertEq(Alice.balance, 10 ether - uint256(totalDistribute)); + for (uint256 i = 0; i < config.length; i++) { + assertEq(config[i].to.balance, config[i].amount); + } + } + + function testFuzz_distributeERC20Fixed(uint8 addressesAmount, uint32 amountToDistribute) + public + { + vm.assume(addressesAmount > 0); + uint240 totalDistribute = uint240(amountToDistribute) * uint240(addressesAmount); + mockERC20.transfer(Alice, totalDistribute); + address[] memory addressList = _generateAccounts(addressesAmount); + (uint8 v, bytes32 r, bytes32 s) = _simplePermit(Alice, totalDistribute); + vm.prank(Alice); + miller.distributeERC20Fixed( + uint240(amountToDistribute), + addressList, + address(mockERC20), + totalDistribute, + 1 days, + v, + r, + s + ); + assertEq(mockERC20.balanceOf(Alice), 0); + for (uint256 i = 0; i < addressList.length; i++) { + assertEq(mockERC20.balanceOf(addressList[i]), amountToDistribute); + } + } + + function testFuzz_distributeERC20(uint8 addressesAmount) public { + vm.assume(addressesAmount > 0); + (Miller.DistributionConfig[] memory config, uint240 totalDistribute) = + _generateConfig(addressesAmount); + mockERC20.transfer(Alice, totalDistribute); + address[] memory addressList = _generateAccounts(addressesAmount); + (uint8 v, bytes32 r, bytes32 s) = _simplePermit(Alice, totalDistribute); + vm.prank(Alice); + miller.distributeERC20(config, address(mockERC20), totalDistribute, 1 days, v, r, s); + assertEq(mockERC20.balanceOf(Alice), 0); + for (uint256 i = 0; i < addressList.length; i++) { + assertEq(mockERC20.balanceOf(config[i].to), config[i].amount); + } + } } diff --git a/test/mocks/MockERC20.sol b/test/mocks/MockERC20.sol new file mode 100644 index 0000000..1213204 --- /dev/null +++ b/test/mocks/MockERC20.sol @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: GPL-3.0 + +pragma solidity ^0.8.16; + +import {ERC20Permit, ERC20} from "@openzeppelin/contracts/token/ERC20/extensions/ERC20Permit.sol"; + +contract MockERC20 is ERC20Permit { + constructor() ERC20Permit("Random") ERC20("MockERC20", "ME20") { + uint256 amount = 1 * 10 ** 9 * 10 ** 18; + _mint(msg.sender, amount); + } +} diff --git a/test/utils/SigUtils.sol b/test/utils/SigUtils.sol new file mode 100644 index 0000000..9ba5bb3 --- /dev/null +++ b/test/utils/SigUtils.sol @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.16; + +contract SigUtils { + bytes32 internal DOMAIN_SEPARATOR; + + constructor(bytes32 _DOMAIN_SEPARATOR) { + DOMAIN_SEPARATOR = _DOMAIN_SEPARATOR; + } + + // keccak256("Permit(address owner,address spender,uint256 value,uint256 nonce,uint256 + // deadline)"); + bytes32 public constant PERMIT_TYPEHASH = + 0x6e71edae12b1b97f4d1f60370fef10105fa2faae0126114a169c64845d6126c9; + + struct Permit { + address owner; + address spender; + uint256 value; + uint256 nonce; + uint256 deadline; + } + + // computes the hash of a permit + function getStructHash(Permit memory _permit) internal pure returns (bytes32) { + return keccak256( + abi.encode( + PERMIT_TYPEHASH, + _permit.owner, + _permit.spender, + _permit.value, + _permit.nonce, + _permit.deadline + ) + ); + } + + // computes the hash of the fully encoded EIP-712 message for the domain, which can be used to + // recover the signer + function getTypedDataHash(Permit memory _permit) public view returns (bytes32) { + return keccak256(abi.encodePacked("\x19\x01", DOMAIN_SEPARATOR, getStructHash(_permit))); + } +}