From 315616a2838c6792db3ef0901bb781a029df563c Mon Sep 17 00:00:00 2001
From: Ana Julia <anajuliabit@gmail.com>
Date: Sun, 30 Jun 2024 17:46:59 -0300
Subject: [PATCH] test constructor

---
 src/Staking.sol               | 6 ++++--
 test/RewardsDistributor.t.sol | 7 +++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/Staking.sol b/src/Staking.sol
index cc7ee4a..256993a 100644
--- a/src/Staking.sol
+++ b/src/Staking.sol
@@ -186,6 +186,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         rewardsDistributor = IRewardsDistributor(_rewardsDistributor);
         lockPeriod = _lockPeriod;
         minStake = _minStake;
+
+        nextStakeId = 1;
     }
 
     /// @notice Stake SHU
@@ -199,7 +201,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /// TODO slippage protection
     function stake(
         uint256 amount
-    ) external onlyKeyper updateRewards returns (uint256) {
+    ) external onlyKeyper updateRewards returns (uint256 stakeId) {
         /////////////////////////// CHECKS ///////////////////////////////
         require(amount > 0, ZeroAmount());
 
@@ -224,7 +226,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         _mint(keyper, sharesToMint);
 
         // Get next stake id and increment it
-        uint256 stakeId = ++nextStakeId;
+        stakeId = nextStakeId++;
 
         stakes[stakeId] = Stake({
             amount: amount,
diff --git a/test/RewardsDistributor.t.sol b/test/RewardsDistributor.t.sol
index 55f1296..cd1108b 100644
--- a/test/RewardsDistributor.t.sol
+++ b/test/RewardsDistributor.t.sol
@@ -35,6 +35,13 @@ contract RewardsDistributorTest is Test {
     }
 }
 
+contract Constructor is RewardsDistributorTest {
+    function test_SetUp() public {
+        assertEq(address(rewardsDistributor.rewardToken()), address(govToken));
+        assertEq(Ownable(address(rewardsDistributor)).owner(), address(this));
+    }
+}
+
 contract OwnableFunctions is RewardsDistributorTest {
     function testFuzz_SetRewardConfigurationEmitEvent(
         address _receiver,