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,