unit tests

dgma · Jul 22, 2024 · b8bdb4d · b8bdb4d
1 parent baf243a
commit b8bdb4d
Show file tree

Hide file tree

Showing 4 changed files with 155 additions and 8 deletions.
diff --git 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
@@ -10,34 +10,126 @@ 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));
         }
         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
@@ -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
@@ -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)));
+    }
+}