fix(tests): fix GaugeController unit tests

pooltogether · May 19, 2022 · 52ce1d6 · 52ce1d6
1 parent f938846
commit 52ce1d6
Show file tree

Hide file tree

Showing 3 changed files with 259 additions and 137 deletions.
diff --git a/contracts/GaugeController.sol b/contracts/GaugeController.sol
@@ -3,13 +3,14 @@
 pragma solidity 0.8.6;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@pooltogether/owner-manager-contracts/contracts/Ownable.sol";
 
 import "./interfaces/IGaugeController.sol";
 import "./interfaces/IGaugeReward.sol";
 import "./libraries/TwabLib.sol";
 import "./libraries/ExtendedSafeCastLib.sol";
 
-contract GaugeController is IGaugeController {
+contract GaugeController is IGaugeController, Ownable {
     using ExtendedSafeCastLib for uint256;
 
     struct GaugeInfo {
@@ -43,11 +44,37 @@ contract GaugeController is IGaugeController {
      */
     mapping(address => TwabLib.Account) internal gaugeScaleTwabs;
 
-    constructor(IERC20 _token, IGaugeReward _gaugeReward) {
+    /* ============ Events ============ */
+
+    /**
+     * @notice Event emitted when the GaugeReward contract address is set
+     * @param gaugeReward Address of the newly set GaugeReward contract
+     */
+    event GaugeRewardSet(IGaugeReward gaugeReward);
+
+    /**
+     * @notice Event emitted when the contract is deployed
+     * @param token Address of the token being staked in the gauge
+     */
+    event Deployed(IERC20 token);
+
+    /* ============ Constructor ============ */
+
+    /**
+     * @notice GaugeController constructor
+     * @param _token Address of the token being staked in the gauge
+     * @param _owner Address of the contract owner
+     */
+    constructor(IERC20 _token, address _owner) Ownable(_owner)  {
+        require(_owner != address(0), "GC/owner-not-zero-address");
+        require(address(_token) != address(0), "GC/token-not-zero-address");
         token = _token;
-        gaugeReward = _gaugeReward;
+
+        emit Deployed(_token);
     }
 
+    /* ============ External Functions ============ */
+
     function deposit(address _to, uint256 _amount) public {
         balances[_to] += _amount;
         token.transferFrom(msg.sender, address(this), _amount);
@@ -92,17 +119,11 @@ contract GaugeController is IGaugeController {
     }
 
     function addGauge(address _gauge) public {
-        addGaugeWithScale(_gauge, 1 ether);
+        _addGaugeWithScale(_gauge, 1 ether);
     }
 
     function addGaugeWithScale(address _gauge, uint256 _scale) public {
-        TwabLib.Account storage gaugeScaleTwab = gaugeScaleTwabs[_gauge];
-        (TwabLib.AccountDetails memory twabDetails, , ) = TwabLib.increaseBalance(
-            gaugeScaleTwab,
-            _scale.toUint208(),
-            uint32(block.timestamp)
-        );
-        gaugeScaleTwab.details = twabDetails;
+        _addGaugeWithScale(_gauge, _scale);
     }
 
     function removeGauge(address _gauge) public {
@@ -117,6 +138,17 @@ contract GaugeController is IGaugeController {
         gaugeScaleTwab.details = twabDetails;
     }
 
+    /**
+     * @notice Set GaugeReward contract
+     * @param _gaugeReward Address of the GaugeReward contract
+     */
+    function setGaugeReward(IGaugeReward _gaugeReward) external onlyOwner {
+        require(address(_gaugeReward) != address(0), "GC/GaugeReward-not-zero-address");
+        gaugeReward = _gaugeReward;
+
+        emit GaugeRewardSet(_gaugeReward);
+    }
+
     function setGaugeScale(address _gauge, uint256 _scale) public {
         TwabLib.Account storage gaugeScaleTwab = gaugeScaleTwabs[_gauge];
         TwabLib.AccountDetails memory twabDetails = gaugeScaleTwab.details;
@@ -184,6 +216,18 @@ contract GaugeController is IGaugeController {
         return userGaugeBalance[_user][_gauge];
     }
 
+    /* ============ Internal Functions ============ */
+
+    function _addGaugeWithScale(address _gauge, uint256 _scale) internal {
+        TwabLib.Account storage gaugeScaleTwab = gaugeScaleTwabs[_gauge];
+        (TwabLib.AccountDetails memory twabDetails, , ) = TwabLib.increaseBalance(
+            gaugeScaleTwab,
+            _scale.toUint208(),
+            uint32(block.timestamp)
+        );
+        gaugeScaleTwab.details = twabDetails;
+    }
+
     function _getAverageGaugeBetween(
         address _gauge,
         uint256 _startTime,

diff --git a/contracts/GaugeReward.sol b/contracts/GaugeReward.sol
@@ -10,8 +10,6 @@ import "./interfaces/IGaugeReward.sol";
 import "./interfaces/IGaugeController.sol";
 import "./interfaces/IPrizePoolLiquidatorListener.sol";
 
-import "hardhat/console.sol";
-
 /**
   * @title  PoolTogether V4 GaugeReward
   * @author PoolTogether Inc Team
@@ -78,7 +76,7 @@ contract GaugeReward is IGaugeReward, IPrizePoolLiquidatorListener, Multicall {
     /* ============ Events ============ */
 
     /**
-     * @notice Emitted when the contract is initialized
+     * @notice Emitted when the contract is deployed
      * @param gaugeController Address of the GaugeController
      * @param vault Address of the Vault
      */
@@ -202,9 +200,11 @@ contract GaugeReward is IGaugeReward, IPrizePoolLiquidatorListener, Multicall {
         uint256 _oldStakeBalance
     ) external override onlyGaugeController {
         RewardToken memory _rewardToken = _claimPastRewards(_gauge, _user, _oldStakeBalance);
+
         if (address(_rewardToken.token) != address(0)) {
             _claim(_gauge, _rewardToken.token, _user, _oldStakeBalance, false);
         }
+
         userLastClaimedTimestamp[_user] = block.timestamp;
     }
 
@@ -331,6 +331,7 @@ contract GaugeReward is IGaugeReward, IPrizePoolLiquidatorListener, Multicall {
 
         if (_gaugeRewardTokensLength > 0) {
             uint256 i = _gaugeRewardTokensLength;
+
             while (i > 0) {
                 i = i - 1;
                 _rewardToken = gaugeRewardTokens[_gauge][i];