Skip to content

Commit

Permalink
updated contracts
Browse files Browse the repository at this point in the history
  • Loading branch information
@gab0071
gab0071 committed Aug 22, 2024
1 parent be7cf3e commit ae0718b
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 40 deletions.
21 changes: 21 additions & 0 deletions layoutContract.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
// Layout of Contract:
// version
// imports
// errors
// interfaces, libraries, contracts
// Type declarations
// State variables
// Events
// Modifiers
// Functions

// Layout of Functions:
// constructor
// receive function (if exists)
// fallback function (if exists)
// external
// public
// internal
// private
// internal & private view & pure functions
// external & public view & pure functions
47 changes: 30 additions & 17 deletions src/Main.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,20 @@ import { RewardToken } from "./RewardToken.sol";
/// @notice This is the main contract that manages the staking and unstaking also the rewardTokensDistribution logic
/// @dev This contract is used just as a educational purpose dont use it in production
contract Main is Ownable {
/*/////////////////////////////////////////////////////////////////////////////////////////////////////////
Custom Errors
/////////////////////////////////////////////////////////////////////////////////////////////////////////*/
error Main__amountMustBeGreaterThanZero();
error Main__stakingBalanceMustBeGreaterThanZero();
error Main__TransferFailed();


/*/////////////////////////////////////////////////////////////////////////////////////////////////////////
Variables
/////////////////////////////////////////////////////////////////////////////////////////////////////////*/
StakingToken public immutable i_stakingToken;
RewardToken public immutable i_rewardToken;
address[] public stakers;
address[] private stakers;

struct StakeInfo {
uint256 StakedBalance;
Expand All @@ -31,12 +39,6 @@ contract Main is Ownable {
event SuccessfulStaked(address indexed user, uint256 amount);
event SuccessfulUnstake(address indexed user, uint256 amount);


/*/////////////////////////////////////////////////////////////////////////////////////////////////////////
Custom Errors
/////////////////////////////////////////////////////////////////////////////////////////////////////////*/
error Main_amountMustBeGreaterThanZero();
error Main_stakingBalanceMustBeGreaterThanZero();


/*/////////////////////////////////////////////////////////////////////////////////////////////////////////
Expand All @@ -57,23 +59,26 @@ contract Main is Ownable {
/// @dev and update the staked balance
/// @dev and lets add to the array the staker address if they are not already there
/// @dev and lets update the staking status
/// @dev if the amount to stake is not good throw an error
/// @dev Remember apply CEI pattern, in this case already applied
function stake(uint256 _amountToStake) public {
if (_amountToStake <= 0) {
revert Main_amountMustBeGreaterThanZero();
revert Main__amountMustBeGreaterThanZero();
}
// If the amount to stake is good lets transfer StakingToken (STK) to the Main contract
i_stakingToken.transferFrom(msg.sender, address(this), _amountToStake);
// update the staked balance
Stakes[msg.sender].StakedBalance += _amountToStake;
// And let's add the staker's address to the array if that user is not already there
if (!Stakes[msg.sender].hasStaked) {
stakers.push(msg.sender);
}
// update the staked balance
Stakes[msg.sender].StakedBalance += _amountToStake;
emit SuccessfulStaked(msg.sender, _amountToStake);
// and lets update the staking status
Stakes[msg.sender].hasStaked = true;
Stakes[msg.sender].isStaking = true;
emit SuccessfulStaked(msg.sender, _amountToStake);
// If the amount to stake is good lets transfer StakingToken (STK) to the Main contract
(bool success) = i_stakingToken.transferFrom(msg.sender, address(this), _amountToStake);
if(!success) {
revert Main__TransferFailed();
}
}

/// @notice This function is used to unstake
Expand All @@ -86,15 +91,15 @@ contract Main is Ownable {
// lets check the balance to unstake
uint256 amountToUnstake = Stakes[msg.sender].StakedBalance;
if (amountToUnstake <= 0) {
revert Main_stakingBalanceMustBeGreaterThanZero();
revert Main__stakingBalanceMustBeGreaterThanZero();
}
// lets transfer to the user the StakingToken (STK) that they staked
i_stakingToken.transfer(msg.sender, amountToUnstake);
// and lets update the staked balance
Stakes[msg.sender].StakedBalance = 0;
// and lets update the staking status
Stakes[msg.sender].isStaking = false;
emit SuccessfulUnstake(msg.sender, amountToUnstake);
// lets transfer to the user the StakingToken (STK) that they staked
i_stakingToken.transfer(msg.sender, amountToUnstake);
}

/// @notice This function is used to distribute the reward token (RWT) to the users who currently staking in the platform
Expand All @@ -112,4 +117,12 @@ contract Main is Ownable {
}
}
}

/*/////////////////////////////////////////////////////////////////////////////////////////////////////////
Getter Functions
/////////////////////////////////////////////////////////////////////////////////////////////////////////*/
function getStakerAt(uint256 index) public view returns (address) {
require(index < stakers.length, "Index out of bounds");
return stakers[index];
}
}
52 changes: 29 additions & 23 deletions test/MainTest.t.sol
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
// SPDX-Lincese-Identifier: MIT
pragma solidity ^0.8.20;

import { StakingToken } from "../src/StakingToken.sol";
import { StakingToken } from "../src/StakingToken.sol";
import { RewardToken } from "../src/RewardToken.sol";
import { Main } from "../src/Main.sol";
import { Test } from "forge-std/Test.sol";
Expand All @@ -17,18 +17,15 @@ contract MainTest is Test {
address public user2 = makeAddr("user2");
address public user3 = makeAddr("user3");

uint256 public stakingTokenInitialSupply = 1_000_000;
uint256 public rewardTokenInitialSupply = 1_000_000;

// custom errors
error Main_amountMustBeGreaterThanZero();
error Main_stakingBalanceMustBeGreaterThanZero();
uint256 public constant TokensInitialSupply = 1_000_000;
uint256 constant amountToStakeUser1 = 555;
uint256 constant amountToStakeUser2 = 4_000;

event SuccessfulUnstake(address indexed user, uint256 amount);

function setUp() public {
stakingToken = new StakingToken(stakingTokenInitialSupply);
rewardToken = new RewardToken(rewardTokenInitialSupply);
stakingToken = new StakingToken(TokensInitialSupply);
rewardToken = new RewardToken(TokensInitialSupply);
mainContract = new Main(stakingToken, rewardToken);

vm.prank(address(this));
Expand All @@ -48,15 +45,16 @@ contract MainTest is Test {
/////////////////////////////////////////////////////////////////////////////////////////////////////////*/
function testStakingToken() public {
vm.startPrank(user1);
uint256 amountToStakeUser1 = 555;
// Let's approve the mainContract contract to stake the STK tokens
stakingToken.approve(address(mainContract), amountToStakeUser1);
// Let's stake the STK tokens
mainContract.stake(amountToStakeUser1);
vm.stopPrank();

uint256 stakingTokenBalanceOfMainContract = stakingToken.balanceOf(address(mainContract));

// Test that the mainContract contract receives the STK tokens
assertEq(stakingToken.balanceOf(address(mainContract)), amountToStakeUser1);
assertEq(stakingTokenBalanceOfMainContract, amountToStakeUser1);

// Destructure the returned struct into individual variables for user1
(uint256 stakedBalance, bool hasStaked, bool isStaking) = mainContract.Stakes(user1);
Expand All @@ -67,24 +65,27 @@ contract MainTest is Test {

// Let's make user2 stake some STK tokens
vm.startPrank(user2);
uint256 amountToStakeUser2 = 4_000;
stakingToken.approve(address(mainContract), amountToStakeUser2);
mainContract.stake(amountToStakeUser2);
vm.stopPrank();

// check the new balance of STK tokens in the mainContract
uint256 newStakedBalance = amountToStakeUser1 + amountToStakeUser2;
assertEq(stakingToken.balanceOf(address(mainContract)), newStakedBalance);
uint256 stakingTokenBalanceOfMainContractNew = stakingToken.balanceOf(address(mainContract));

assertEq(stakingTokenBalanceOfMainContractNew, newStakedBalance);

// verify that user2 was succefully added into the stakers array
address staker = mainContract.stakers(1);
address staker = mainContract.getStakerAt(1);
assertEq(staker, user2, "Staker number 2 was not added to the array.");



// check that user3 can not stake with 0 amount
vm.startPrank(user3);
uint256 amountToStakeUser3 = 0;
stakingToken.approve(address(mainContract), amountToStakeUser3);
vm.expectRevert(Main_amountMustBeGreaterThanZero.selector);
vm.expectRevert(Main.Main__amountMustBeGreaterThanZero.selector);
mainContract.stake(amountToStakeUser3);
vm.stopPrank();
}
Expand All @@ -94,9 +95,10 @@ contract MainTest is Test {
/////////////////////////////////////////////////////////////////////////////////////////////////////////*/

function testUnstake() public {
uint256 amountToUnstake = 555;
// Let's call the function where users stake tokens
testStakingToken();
// Let's get the staked balance of user1
(uint256 amountToUnstake, , ) = mainContract.Stakes(user1);

// Let's emit the event where the user unstakes
vm.expectEmit();
Expand All @@ -109,7 +111,7 @@ contract MainTest is Test {

// Let's see if the user3 can unstake
vm.startPrank(user3);
vm.expectRevert(Main_stakingBalanceMustBeGreaterThanZero.selector);
vm.expectRevert(Main.Main__stakingBalanceMustBeGreaterThanZero.selector);
mainContract.unstake();
vm.stopPrank();
}
Expand All @@ -120,20 +122,24 @@ contract MainTest is Test {
function testRewardTokensDistribution() public {
// Let's create the enviroment where exist stakers
testStakingToken();

uint256 amountUser1Staked = 555;
uint256 amountUser2Staked = 4_000;

// Access to the struct where exist the staked balance of the stakers
(uint256 amountUser1Staked, , ) = mainContract.Stakes(user1);
(uint256 amountUser2Staked, , ) = mainContract.Stakes(user2);

// Now lets make the distribution of the reward tokens to the stakers (In this case 1:1)
// So if user1 staked 555 STK tokens he is going to receive 555 RWT tokens
vm.startPrank(address(this));
// We are going to send full supply to the mainContract contract to distribute the reward tokens to the stakers
rewardToken.transfer(address(mainContract), rewardTokenInitialSupply);
rewardToken.transfer(address(mainContract), TokensInitialSupply);
mainContract.rewardTokensDistribution();
vm.stopPrank();

uint256 rewardTokenBalanceOfUser1 = rewardToken.balanceOf(address(user1));
uint256 rewardTokenBalanceOfUser2 = rewardToken.balanceOf(address(user2));

// Let's check if the reward token balance is correct
assertEq(rewardToken.balanceOf(user1), amountUser1Staked);
assertEq(rewardToken.balanceOf(user2), amountUser2Staked);
assertEq(rewardTokenBalanceOfUser1, amountUser1Staked);
assertEq(rewardTokenBalanceOfUser2, amountUser2Staked);
}
}

0 comments on commit ae0718b

Please sign in to comment.