invariant tests

src/test/invariant/ERC20GaugesInvariantTest.t.sol

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+pragma solidity 0.8.10;
+
+import {DSTestPlus} from "solmate/test/utils/DSTestPlus.sol";
+import {MockERC20Gauges} from "../mocks/MockERC20Gauges.sol";
+
+contract ERC20GaugesInvariantTest is DSTestPlus {
+
+    MockERC20Gauges token;
+
+    function setUp() public {
+        token = new MockERC20Gauges(address(this), 3600); // 1 hour cycles
+        token.mint(address(this), 100e18);
+    }
+
+    function invariant_userWeight() public {
+        require(token.getUserWeight(address(this)) <= token.balanceOf(address(this)));
+        require(token.userUnusedWeight(address(this)) == token.balanceOf(address(this)) - token.getUserWeight(address(this)));
+        require(token.userWeightSum(address(this)) == token.getUserWeight(address(this)));
+    }
+
+    function invariant_maxGauges() public {
+        require(token.canContractExceedMaxGauges(address(this)) || token.numUserGauges(address(this)) <= token.maxGauges());
+    }
+
+    function invariant_currentCycle() public {
+        require(token.getCurrentCycle() >= block.timestamp);
+        require(token.getCurrentCycle() % token.gaugeCycleLength() == 0);
+    }
+
+    function invariant_totalWeight() public {
+        require(token.totalWeight() <= token.totalSupply());
+        require(token.totalWeight() == token.gaugeWeightSum());
+    }
+
+    function invariant_storedTotalWeight() public {
+        require(token.storedTotalWeight() == token.storedGaugeWeightSum());
+    }
+
+    function invariant_allocation() public {
+        require(token.summedGaugeAllocation(1_000_000e18) <= 1_000_000e18);
+    }
+}

src/test/invariant/ERC20MultiVotesInvariantTest.t.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+pragma solidity 0.8.10;
+
+import {DSTestPlus} from "solmate/test/utils/DSTestPlus.sol";
+import {MockERC20MultiVotes, ERC20MultiVotes} from "../mocks/MockERC20MultiVotes.sol";
+
+contract ERC20MultiVotesTest is DSTestPlus {
+
+    MockERC20MultiVotes token;
+    address constant delegate1 = address(0xDEAD);
+    address constant delegate2 = address(0xBEEF);
+
+    function setUp() public {
+        token = new MockERC20MultiVotes(address(this));
+        token.mint(address(this), 100e18);
+        token.setMaxDelegates(2);
+    }
+
+    function invariant_userVotes() public {
+        require(token.userDelegatedVotes(address(this)) <= token.balanceOf(address(this)));
+        require(token.freeVotes(address(this)) == token.balanceOf(address(this)) - token.userDelegatedVotes(address(this)));
+        require(token.userDelegateSum(address(this)) == token.userDelegatedVotes(address(this)));
+    }
+
+    function invariant_maxDelegates() public {
+        require(token.canContractExceedMaxDelegates(address(this)) || token.delegateCount(address(this)) <= token.maxDelegates());
+    }
+}

src/test/mocks/MockERC20Gauges.sol

@@ -2,9 +2,15 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 pragma solidity >=0.8.0;
 
-import {ERC20Gauges, ERC20, Auth, Authority} from "../../token/ERC20Gauges.sol";
+import "../../token/ERC20Gauges.sol";
+
+import {Hevm} from "solmate/test/utils/Hevm.sol";
 
 contract MockERC20Gauges is ERC20Gauges {
+    using EnumerableSet for EnumerableSet.AddressSet;
+
+    Hevm internal constant hevm = Hevm(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D);
+
     constructor(
         address _owner,
         uint32 _cycleLength
@@ -17,4 +23,42 @@ contract MockERC20Gauges is ERC20Gauges {
     function burn(address from, uint256 value) public virtual {
         _burn(from, value);
     }
+
+    ////// Invariant test helpers
+
+    function incrementGaugeByNum(uint256 gaugeNum, uint112 weight) public virtual {
+        uint256 max = _gauges.length();
+        gaugeNum = gaugeNum % max;
+        uint32 currentCycle = getCurrentCycle();
+        _incrementGaugeWeight(msg.sender, _gauges.at(gaugeNum), weight, currentCycle);
+        _incrementUserAndGlobalWeights(msg.sender, weight, currentCycle);
+    }
+
+    function warpCycle(uint256 warp) public {
+        hevm.warp(block.timestamp + (warp % gaugeCycleLength));
+    }
+
+    function gaugeWeightSum() public view virtual returns (uint112 sum) {
+        for (uint256 i = 0; i < _gauges.length(); i++) {
+            sum += getGaugeWeight(_gauges.at(i));
+        }
+    }
+
+    function storedGaugeWeightSum() public view virtual returns (uint112 sum) {
+        for (uint256 i = 0; i < _gauges.length(); i++) {
+            sum += getStoredGaugeWeight(_gauges.at(i));
+        }
+    }
+
+    function userWeightSum(address user) public view returns (uint112 sum) {
+        for (uint256 i = 0; i < _userGauges[user].length(); i++) {
+            sum += getGaugeWeight(_userGauges[user].at(i));
+        }
+    }
+
+    function summedGaugeAllocation(uint256 quantity) public view returns(uint256 sum) {
+        for (uint256 i = 0; i < _gauges.length(); i++) {
+            sum += this.calculateGaugeAllocation(_gauges.at(i), quantity);
+        }
+    }
 }

src/test/mocks/MockERC20MultiVotes.sol

@@ -2,9 +2,11 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 pragma solidity >=0.8.0;
 
-import {ERC20MultiVotes, ERC20, Auth, Authority} from "../../token/ERC20MultiVotes.sol";
+import "../../token/ERC20MultiVotes.sol";
 
 contract MockERC20MultiVotes is ERC20MultiVotes {
+    using EnumerableSet for EnumerableSet.AddressSet;
+
     constructor(
         address _owner
     ) ERC20("Token", "TKN", 18) Auth(_owner, Authority(address(0))) {}
@@ -16,4 +18,10 @@ contract MockERC20MultiVotes is ERC20MultiVotes {
     function burn(address from, uint256 value) public virtual {
         _burn(from, value);
     }
+
+    function userDelegateSum(address user) public view virtual returns (uint256 sum) {
+        for (uint256 i = 0; i < delegateCount(user); i++) {
+            sum += delegatesVotesCount(user, _delegates[user].at(i));
+        }
+    }
 }

src/token/ERC20MultiVotes.sol

@@ -146,13 +146,13 @@ abstract contract ERC20MultiVotes is ERC20, Auth {
     error DelegationError();
 
     /// @notice mapping from a delegator and delegatee to the delegated amount.
-    mapping(address => mapping(address => uint256)) private _delegatesVotesCount;
+    mapping(address => mapping(address => uint256)) internal _delegatesVotesCount;
 
     /// @notice mapping from a delegator to the total number of delegated votes.
     mapping(address => uint256) public userDelegatedVotes;
 
     /// @notice list of delegates per user.
-    mapping(address => EnumerableSet.AddressSet) private _delegates;
+    mapping(address => EnumerableSet.AddressSet) internal _delegates;
 
     /**
      * @notice Get the amount of votes currently delegated by `delegator` to `delegatee`.