From 0e63a9630fc58349c89b4162c6c158d9aca96520 Mon Sep 17 00:00:00 2001
From: Ana Julia <anajuliabit@gmail.com>
Date: Sat, 15 Jun 2024 16:58:49 -0300
Subject: [PATCH] testAddKeyper

---
 src/RewardsDistributor.sol      |  2 +-
 src/Staking.sol                 | 10 ++++++--
 src/interfaces/IStaking.sol     |  1 +
 test/unit/StakingUnitTest.t.sol | 44 +++++++++++++++++++++++++++------
 4 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/src/RewardsDistributor.sol b/src/RewardsDistributor.sol
index b969875..cc110cd 100644
--- a/src/RewardsDistributor.sol
+++ b/src/RewardsDistributor.sol
@@ -47,7 +47,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
         __Ownable2Step_init();
 
         // Transfer ownership to the DAO contract
-        transferOwnership(newOwner);
+        _transferOwnership(newOwner);
     }
 
     /// @notice Add a reward configuration
diff --git a/src/Staking.sol b/src/Staking.sol
index 6c6336a..cb1593f 100644
--- a/src/Staking.sol
+++ b/src/Staking.sol
@@ -14,7 +14,8 @@ interface IRewardsDistributor {
 }
 
 // TODO should be pausable?
-// TODO use SafeTransferLib to every calculation
+// TODO is this vulnerable to first deposit attack?
+// TODO check calculations
 contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /*//////////////////////////////////////////////////////////////
                                LIBRARIES
@@ -86,6 +87,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     );
     event Unstaked(address user, uint256 amount, uint256 shares);
     event ClaimRewards(address user, address rewardToken, uint256 rewards);
+    event KeyperSet(address keyper, bool isKeyper);
 
     /*//////////////////////////////////////////////////////////////
                                  MODIFIERS
@@ -121,7 +123,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         __Ownable2Step_init();
 
         // Transfer ownership to the DAO contract
-        transferOwnership(newOwner);
+        _transferOwnership(newOwner);
 
         stakingToken = IERC20(_stakingToken);
         rewardsDistributor = IRewardsDistributor(_rewardsDistributor);
@@ -408,6 +410,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /// @param isKeyper Whether the keyper is a keyper or not
     function setKeyper(address keyper, bool isKeyper) external onlyOwner {
         keypers[keyper] = isKeyper;
+
+        emit KeyperSet(keyper, isKeyper);
     }
 
     /// @notice Set multiple keypers
@@ -419,6 +423,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     ) external onlyOwner {
         for (uint256 i = 0; i < _keypers.length; i++) {
             keypers[_keypers[i]] = isKeyper;
+
+            emit KeyperSet(_keypers[i], isKeyper);
         }
     }
 
diff --git a/src/interfaces/IStaking.sol b/src/interfaces/IStaking.sol
index 2491267..dcf71a2 100644
--- a/src/interfaces/IStaking.sol
+++ b/src/interfaces/IStaking.sol
@@ -59,4 +59,5 @@ interface IStaking {
     );
     event Unstaked(address user, uint256 amount, uint256 shares);
     event ClaimRewards(address user, address rewardToken, uint256 rewards);
+    event KeyperSet(address keyper, bool isKeyper);
 }
diff --git a/test/unit/StakingUnitTest.t.sol b/test/unit/StakingUnitTest.t.sol
index ea113c3..0a3a1db 100644
--- a/test/unit/StakingUnitTest.t.sol
+++ b/test/unit/StakingUnitTest.t.sol
@@ -13,20 +13,33 @@ import {IRewardsDistributor} from "../../src/interfaces/IRewardsDistributor.sol"
 import {MockShu} from "../mocks/MockShu.sol";
 
 contract StakingUnitTest is Test {
-    IStaking staking;
+    IStaking public staking;
+
+    uint256 constant lockPeriod = 60 * 24 * 30 * 6; // 6 months
+    uint256 constant minStake = 50_000 * 1e18; // 50k
+
+    address keyper = address(0x1234);
 
     function setUp() public {
         // deploy mock shu
         MockShu shu = new MockShu();
 
+        shu.mint(keyper, 1_000_000 * 1e18);
+
         // deploy rewards distributor
         address rewardsDistributor = address(new RewardsDistributor());
 
-        TransparentUpgradeableProxy rewardsDistributionProxy = new TransparentUpgradeableProxy(
+        address rewardsDistributionProxy = address(
+            new TransparentUpgradeableProxy(
                 rewardsDistributor,
                 address(this),
-                abi.encodeWithSignature("initialize(address)", address(this))
-            );
+                ""
+            )
+        );
+
+        RewardsDistributor(rewardsDistributionProxy).initialize(
+            address(this) // owner
+        );
 
         // deploy staking
         address stakingContract = address(new Staking());
@@ -35,9 +48,6 @@ contract StakingUnitTest is Test {
             new TransparentUpgradeableProxy(stakingContract, address(this), "")
         );
 
-        uint256 lockPeriod = 60 * 24 * 30 * 6; // 6 months
-        uint256 minStake = 50_000 * 1e18; // 50k
-
         Staking(address(stakingProxy)).initialize(
             address(this), // owner
             address(shu),
@@ -48,4 +58,24 @@ contract StakingUnitTest is Test {
 
         staking = IStaking(stakingProxy);
     }
+
+    function testAddKeyper() public {
+        vm.expectEmit(address(staking));
+        emit IStaking.KeyperSet(keyper, true);
+        staking.setKeyper(keyper, true);
+    }
+
+    function testStakeSucceed() public {
+        testAddKeyper();
+
+        vm.expectEmit(true, true, true, true, address(staking));
+        emit IStaking.Staked(
+            address(this),
+            minStake,
+            minStake, // first stake, shares == amount
+            lockPeriod
+        );
+        vm.prank(keyper);
+        staking.stake(50_000 * 1e18);
+    }
 }