diff --git a/remappings.txt b/remappings.txt
index 707e5d0..63bbb7d 100644
--- a/remappings.txt
+++ b/remappings.txt
@@ -1,3 +1,4 @@
+@forge-std/=lib/forge-std/src/
 @openzeppelin=lib/openzeppelin-contracts/
 @openzeppelin-upgradeable=lib/openzeppelin-contracts-upgradeable/
 @solmate=lib/solmate/src/
diff --git a/src/RewardsDistribution.sol b/src/RewardsDistributor.sol
similarity index 99%
rename from src/RewardsDistribution.sol
rename to src/RewardsDistributor.sol
index cb8f10f..b969875 100644
--- a/src/RewardsDistribution.sol
+++ b/src/RewardsDistributor.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity 0.8.20;
+pragma solidity 0.8.25;
 
 import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
diff --git a/src/Staking.sol b/src/Staking.sol
index dcca07e..6c6336a 100644
--- a/src/Staking.sol
+++ b/src/Staking.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity 0.8.20;
+pragma solidity 0.8.25;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
@@ -111,8 +111,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /// @param _minStake The minimum stake amount
     function initialize(
         address newOwner,
-        IERC20 _stakingToken,
-        IRewardsDistributor _rewardsDistributor,
+        address _stakingToken,
+        address _rewardsDistributor,
         uint256 _lockPeriod,
         uint256 _minStake
     ) public initializer {
@@ -123,8 +123,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         // Transfer ownership to the DAO contract
         transferOwnership(newOwner);
 
-        stakingToken = _stakingToken;
-        rewardsDistributor = _rewardsDistributor;
+        stakingToken = IERC20(_stakingToken);
+        rewardsDistributor = IRewardsDistributor(_rewardsDistributor);
         lockPeriod = _lockPeriod;
         minStake = _minStake;
     }
diff --git a/src/interfaces/IRewardsDistributor.sol b/src/interfaces/IRewardsDistributor.sol
new file mode 100644
index 0000000..f980966
--- /dev/null
+++ b/src/interfaces/IRewardsDistributor.sol
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.25;
+
+interface IRewardsDistributor {
+    struct RewardConfiguration {
+        address token; // the reward token
+        uint256 emissionRate; // emission per second
+        uint256 lastUpdate; // last update timestamp
+    }
+
+    function addRewardConfiguration(
+        address receiver,
+        address token,
+        uint256 emissionRate
+    ) external;
+
+    function updateEmissonRate(
+        address receiver,
+        address token,
+        uint256 emissionRate
+    ) external;
+
+    function distributeReward(address token) external;
+
+    function rewardConfigurations(
+        address receiver,
+        uint256 index
+    )
+        external
+        view
+        returns (address token, uint256 emissionRate, uint256 lastUpdate);
+
+    function rewardConfigurationsIds(
+        address receiver,
+        address token
+    ) external view returns (uint256);
+}
diff --git a/src/interfaces/IStaking.sol b/src/interfaces/IStaking.sol
new file mode 100644
index 0000000..2491267
--- /dev/null
+++ b/src/interfaces/IStaking.sol
@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.25;
+
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {IRewardsDistributor} from "./IRewardsDistributor.sol";
+
+interface IStaking {
+    function initialize(
+        address newOwner,
+        IERC20 _stakingToken,
+        IRewardsDistributor _rewardsDistributor,
+        uint256 _lockPeriod,
+        uint256 _minStake
+    ) external;
+
+    function stake(uint256 amount) external;
+
+    function unstake(
+        address keyper,
+        uint256 stakeIndex,
+        uint256 amount
+    ) external;
+
+    function claimReward(IERC20 rewardToken, uint256 amount) external;
+
+    function harvest(address keyper) external;
+
+    function setRewardsDistributor(
+        IRewardsDistributor _rewardsDistributor
+    ) external;
+
+    function setLockPeriod(uint256 _lockPeriod) external;
+
+    function setMinStake(uint256 _minStake) external;
+
+    function setKeyper(address keyper, bool isKeyper) external;
+
+    function setKeypers(address[] memory _keypers, bool isKeyper) external;
+
+    function addRewardToken(address rewardToken) external;
+
+    function convertToShares(uint256 assets) external view returns (uint256);
+
+    function convertToAssets(uint256 shares) external view returns (uint256);
+
+    function totalAssets() external view returns (uint256);
+
+    function maxWithdraw(address keyper) external view returns (uint256);
+
+    function maxClaimableRewards(
+        address keyper
+    ) external view returns (uint256);
+
+    event Staked(
+        address indexed user,
+        uint256 indexed amount,
+        uint256 indexed shares,
+        uint256 lockPeriod
+    );
+    event Unstaked(address user, uint256 amount, uint256 shares);
+    event ClaimRewards(address user, address rewardToken, uint256 rewards);
+}
diff --git a/test/mocks/MockShu.sol b/test/mocks/MockShu.sol
new file mode 100644
index 0000000..3e31c66
--- /dev/null
+++ b/test/mocks/MockShu.sol
@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.25;
+
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
+contract MockShu is ERC20 {
+    constructor() ERC20("Shu", "SHU") {
+        _mint(msg.sender, 5_000_000 * 1e18);
+    }
+
+    function mint(address account, uint256 amount) external {
+        _mint(account, amount);
+    }
+}
diff --git a/test/unit/StakingUnitTest.t.sol b/test/unit/StakingUnitTest.t.sol
new file mode 100644
index 0000000..ea113c3
--- /dev/null
+++ b/test/unit/StakingUnitTest.t.sol
@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.25;
+
+import "@forge-std/Test.sol";
+
+import {TransparentUpgradeableProxy, ITransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+
+import {Staking} from "../../src/Staking.sol";
+import {IStaking} from "../../src/interfaces/IStaking.sol";
+import {RewardsDistributor} from "../../src/RewardsDistributor.sol";
+import {IRewardsDistributor} from "../../src/interfaces/IRewardsDistributor.sol";
+
+import {MockShu} from "../mocks/MockShu.sol";
+
+contract StakingUnitTest is Test {
+    IStaking staking;
+
+    function setUp() public {
+        // deploy mock shu
+        MockShu shu = new MockShu();
+
+        // deploy rewards distributor
+        address rewardsDistributor = address(new RewardsDistributor());
+
+        TransparentUpgradeableProxy rewardsDistributionProxy = new TransparentUpgradeableProxy(
+                rewardsDistributor,
+                address(this),
+                abi.encodeWithSignature("initialize(address)", address(this))
+            );
+
+        // deploy staking
+        address stakingContract = address(new Staking());
+
+        address stakingProxy = address(
+            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),
+            address(rewardsDistributionProxy),
+            lockPeriod,
+            minStake
+        );
+
+        staking = IStaking(stakingProxy);
+    }
+}