setRewardConfiguration

src/RewardsDistributor.sol

@@ -3,7 +3,6 @@ pragma solidity 0.8.26;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import {SafeTransferLib} from "@solmate/utils/SafeTransferLib.sol";
 import {Ownable2StepUpgradeable} from "@openzeppelin-upgradeable/contracts/access/Ownable2StepUpgradeable.sol";
 import {IRewardsDistributor} from "./interfaces/IRewardsDistributor.sol";
 
@@ -49,6 +48,15 @@ contract RewardsDistributor is Ownable2StepUpgradeable, IRewardsDistributor {
 
     event RewardCollected(address indexed receiver, uint256 reward);
 
+    event RewardTokenSet(address indexed rewardToken);
+
+    /*//////////////////////////////////////////////////////////////
+                                 ERRORS
+    //////////////////////////////////////////////////////////////*/
+
+    /// @notice Thrown when address is zero
+    error ZeroAddress();
+
     /// @notice Initialize the contract
     /// @param newOwner The owner of the contract, i.e. the DAO contract address
     constructor(address newOwner, address _rewardToken) {
@@ -93,7 +101,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable, IRewardsDistributor {
         address receiver,
         uint256 emissionRate
     ) external override onlyOwner {
-        require(receiver != address(0), "Invalid receiver");
+        require(receiver != address(0), ZeroAddress());
 
         rewardConfigurations[receiver] = RewardConfiguration(
             emissionRate,
@@ -103,10 +111,25 @@ contract RewardsDistributor is Ownable2StepUpgradeable, IRewardsDistributor {
         emit RewardConfigurationSet(receiver, emissionRate);
     }
 
+    /// @notice Withdraw funds from the contract
+    /// @param to The address to withdraw to
+    /// @param amount The amount to withdraw
     function withdrawFunds(
         address to,
         uint256 amount
-    ) external override onlyOwner {
+    ) public override onlyOwner {
         rewardToken.safeTransfer(to, amount);
     }
+
+    /// @notice Set the reward token
+    /// @param _rewardToken The reward token
+    function setRewardToken(address _rewardToken) external onlyOwner {
+        // withdraw remaining old reward token
+        withdrawFunds(msg.sender, rewardToken.balanceOf(address(this)));
+
+        // set the new reward token
+        rewardToken = IERC20(_rewardToken);
+
+        emit RewardTokenSet(_rewardToken);
+    }
 }

src/Staking.sol

@@ -119,7 +119,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /// amount is less than the minimum stake set by the DAO
     error FirstStakeLessThanMinStake();
 
-    /// @notice Trownn when amount is zero
+    /// @notice Trown when amount is zero
     error ZeroAmount();
 
     /// @notice Thrown when someone try to unstake a stake that doesn't belong

src/interfaces/IRewardsDistributor.sol

@@ -10,4 +10,6 @@ interface IRewardsDistributor {
         address receiver,
         uint256 emissionRate
     ) external;
+
+    function setRewardToken(address _rewardToken) external;
 }

test/RewardsDistributor.t.sol

@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.26;
+
+import "@forge-std/Test.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+
+import {RewardsDistributor} from "src/RewardsDistributor.sol";
+import {MockGovToken} from "test/mocks/MockGovToken.sol";
+
+contract RewardsDistributorTest is Test {
+    RewardsDistributor public rewardsDistributor;
+    MockGovToken public govToken;
+
+    function setUp() public {
+        // Set the block timestamp to an arbitrary value to avoid introducing assumptions into tests
+        // based on a starting timestamp of 0, which is the default.
+        _jumpAhead(1234);
+
+        govToken = new MockGovToken();
+        govToken.mint(address(this), 1_000_000_000e18);
+        vm.label(address(govToken), "govToken");
+
+        // deploy rewards distributor
+        rewardsDistributor = new RewardsDistributor(
+            address(this),
+            address(govToken)
+        );
+    }
+
+    function _jumpAhead(uint256 _seconds) public {
+        vm.warp(vm.getBlockTimestamp() + _seconds);
+    }
+
+    function _setRewardConfiguration(
+        address receiver,
+        uint256 emissionRate
+    ) internal {
+        emissionRate = bound(emissionRate, 0, 1e18);
+        rewardsDistributor.setRewardConfiguration(receiver, emissionRate);
+    }
+}
+
+contract OwnableFunctions is RewardsDistributorTest {
+    function test_SetRewardConfigurationEmitEvent(
+        address _receiver,
+        uint256 _emissionRate
+    ) public {
+        vm.expectEmit();
+        emit RewardsDistributor.RewardConfigurationSet(
+            _receiver,
+            _emissionRate
+        );
+        rewardsDistributor.setRewardConfiguration(_receiver, _emissionRate);
+    }
+
+    function test_SetRewardConfigurationSetEmissionRate(
+        address _receiver,
+        uint256 _emissionRate
+    ) public {
+        rewardsDistributor.setRewardConfiguration(_receiver, _emissionRate);
+
+        (uint256 emissionRate, ) = rewardsDistributor.rewardConfigurations(
+            _receiver
+        );
+
+        assertEq(emissionRate, _emissionRate);
+    }
+
+    function test_SetRewardConfigurationSetLastUpdate(
+        address _receiver,
+        uint256 _emissionRate
+    ) public {
+        rewardsDistributor.setRewardConfiguration(_receiver, _emissionRate);
+
+        (, uint256 lastUpdate) = rewardsDistributor.rewardConfigurations(
+            _receiver
+        );
+        assertEq(lastUpdate, vm.getBlockTimestamp());
+    }
+
+    function test_RevertIf_SetRewardConfigurationZeroAddress(
+        uint256 _emissionRate
+    ) public {
+        vm.expectRevert(RewardsDistributor.ZeroAddress.selector);
+        rewardsDistributor.setRewardConfiguration(address(0), _emissionRate);
+    }
+
+    function test_RevertIf_SetRewardConfigurationNotOwner(
+        address _anyone,
+        address _receiver,
+        uint256 _emissionRate
+    ) public {
+        vm.assume(_anyone != address(this));
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                Ownable.OwnableUnauthorizedAccount.selector,
+                _anyone
+            )
+        );
+        vm.prank(_anyone);
+        rewardsDistributor.setRewardConfiguration(_receiver, _emissionRate);
+    }
+}

test/Staking.t.sol

@@ -1394,6 +1394,8 @@ contract OwnableFunctions is StakingTest {
         uint256[] memory keyperStakeIds = staking.getKeyperStakeIds(_keyper);
         assertEq(keyperStakeIds.length, 0, "Wrong stake ids");
     }
+
+    // TEST CASES FOR NON OWNERS
 }
 
 contract ViewFunctions is StakingTest {