updated staking to have virtual methods

deploy/agent-staking-base.sh

@@ -22,4 +22,4 @@ else
         --etherscan-api-key $BASESCAN_API_KEY \
         --account $FORGE_ACCOUNT \
         --slow
-fis
+fi

src/AgentStaking.sol

@@ -17,21 +17,20 @@ contract AgentStaking is OwnableUpgradeable, UUPSUpgradeable {
     error NoLockedWithdrawalsFound();
 
     IERC20 public agentToken;
-    uint256 public constant UNLOCK_TIME = 1 days;
 
     struct LockedWithdrawal {
         uint256 amount;
         uint256 lockedUntil;
     }
 
-    mapping(address => uint256) private stakes;
+    mapping(address => uint256) internal stakes;
 
     // Queue of withdrawals
     // Withdrawals are pushed to the end of the queue
     // withdrawalQueueStartIndexes is used to track the start of the queue since were using an array
     // This is to avoid shifting the array every time a withdrawal is claimed
-    mapping(address => LockedWithdrawal[]) private withdrawalQueue;
-    mapping(address => uint256) private withdrawalQueueStartIndexes;
+    mapping(address => LockedWithdrawal[]) internal withdrawalQueue;
+    mapping(address => uint256) internal withdrawalQueueStartIndexes;
 
     event Stake(address indexed account, uint256 amount, uint256 totalStaked);
     event Unstake(address indexed account, uint256 amount, uint256 unlocksAt, uint256 totalStaked);
@@ -51,7 +50,7 @@ contract AgentStaking is OwnableUpgradeable, UUPSUpgradeable {
 
     /// @notice Stake agent tokens
     /// @param amount The amount of tokens to stake
-    function stake(uint256 amount) external {
+    function stake(uint256 amount) public virtual {
         if (amount == 0) {
             revert EmptyAmount();
         }
@@ -66,7 +65,7 @@ contract AgentStaking is OwnableUpgradeable, UUPSUpgradeable {
     /// @notice Unstake agent tokens
     /// @param amount The amount of tokens to unstake
     /// @dev Tokens will be locked for a period of time (1 day) before they can be claimed
-    function unstake(uint256 amount) external {
+    function unstake(uint256 amount) public virtual {
         if (amount == 0) {
             revert EmptyAmount();
         }
@@ -78,15 +77,15 @@ contract AgentStaking is OwnableUpgradeable, UUPSUpgradeable {
         uint256 totalStaked = stakes[msg.sender] - amount; 
         stakes[msg.sender] = totalStaked;
 
-        uint256 unlocksAt = block.timestamp + UNLOCK_TIME;
+        uint256 unlocksAt = block.timestamp + unlock_time();
 
         withdrawalQueue[msg.sender].push(LockedWithdrawal(amount, unlocksAt));
         emit Unstake(msg.sender, amount, unlocksAt, totalStaked);
     }
 
     /// @notice Claim unlocked agent tokens
     /// @param count The number of 'unlocked' withdrawals to claim
-    function claim(uint256 count, address recipient) external {
+    function claim(uint256 count, address recipient) public virtual {
         uint256 start = withdrawalQueueStartIndexes[msg.sender];
 
         uint256 length = withdrawalQueue[msg.sender].length;
@@ -121,15 +120,15 @@ contract AgentStaking is OwnableUpgradeable, UUPSUpgradeable {
 
     /// @notice Get the amount of agent tokens staked by an account
     /// @param account The account to get the staked amount for
-    function getStakedAmount(address account) external view returns (uint256) {
+    function getStakedAmount(address account) public virtual view returns (uint256) {
         return stakes[account];
     }
 
     /// @notice Get the locked withdrawals for an account
     /// @param account The account to get the withdrawals for
     /// @param start The start index of the withdrawals
     /// @param count The number of withdrawals to get
-    function getWithdrawals(address account, uint256 start, uint256 count) external view returns (LockedWithdrawal[] memory) {
+    function getWithdrawals(address account, uint256 start, uint256 count) public virtual view returns (LockedWithdrawal[] memory) {
         start = withdrawalQueueStartIndexes[account] + start;
         uint256 length = withdrawalQueue[account].length;
 
@@ -152,10 +151,15 @@ contract AgentStaking is OwnableUpgradeable, UUPSUpgradeable {
 
     /// @notice Get the number of locked withdrawals for an account
     /// @param account The account to get the number of withdrawals for
-    function getWithdrawalCount(address account) external view returns (uint256) {
+    function getWithdrawalCount(address account) public virtual view returns (uint256) {
         return withdrawalQueue[account].length - withdrawalQueueStartIndexes[account];
     }
 
+    /// @notice Get the unlock time
+    function unlock_time() public virtual view returns (uint256) {
+        return 1 days;
+    }
+
     /// @dev Only the owner can upgrade the contract
-    function _authorizeUpgrade(address) internal override onlyOwner {}
+    function _authorizeUpgrade(address) internal virtual override onlyOwner {}
 }

test/AgentStaking.t.sol

@@ -4,6 +4,7 @@ pragma solidity 0.8.28;
 import {Test, console} from "forge-std/Test.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 
 import {AgentToken} from "../src/AgentToken.sol";
@@ -424,6 +425,107 @@ contract AgentStakingTest is Test {
         staking.upgradeToAndCall(newImplementation, "");
     }
 
+    function test_upgradeUnlockTime() public {
+        uint256 amount = 100 * 1e18;
+
+        vm.prank(owner);
+        token.transfer(user, amount);
+
+        vm.startPrank(user);
+        token.approve(address(staking), amount);
+        staking.stake(amount);
+
+        address newImplementation = address(new AgentStakingUnlock());
+
+        vm.startPrank(owner);
+        staking.upgradeToAndCall(newImplementation, "");
+
+        vm.startPrank(user);
+        staking.unstake(amount);
+
+        vm.warp(block.timestamp + 1 days);
+
+        staking.claim(1, user);
+
+        assertEq(staking.getStakedAmount(user), 0);
+        assertEq(token.balanceOf(user), 0);
+
+        vm.warp(block.timestamp + 1 days);
+
+        staking.claim(1, user);
+
+        assertEq(staking.getStakedAmount(user), 0);
+        assertEq(token.balanceOf(user), amount);
+    }
+
+    function test_canUpgradeDisableStaking() public {
+        uint256 amount = 100 * 1e18;
+
+        vm.prank(owner);
+        token.transfer(user, amount);
+
+        vm.startPrank(user);
+        token.approve(address(staking), amount);
+        staking.stake(amount);
+
+        address newImplementation = address(new AgentStakingDisabled());
+
+        vm.startPrank(owner);
+        staking.upgradeToAndCall(newImplementation, "");
+
+        vm.startPrank(user);
+        vm.expectRevert("Staking is disabled");
+        staking.stake(amount);
+    }
+
+    function test_canUpgradeDisableUnstaking() public {
+        uint256 amount = 100 * 1e18;
+
+        vm.prank(owner);
+        token.transfer(user, amount);
+
+        vm.startPrank(user);
+        token.approve(address(staking), amount);
+        staking.stake(amount);
+
+        address newImplementation = address(new AgentUnstaking());
+
+        vm.startPrank(owner);
+        staking.upgradeToAndCall(newImplementation, "");
+
+        vm.startPrank(user);
+        vm.expectRevert("Unstaking is disabled");
+        staking.unstake(amount);
+    }
+
+    function test_canUpgradeAndAccessStorage() public {
+        uint256 amount = 100 * 1e18;
+        address newStaking = makeAddr("newStaking");
+
+        vm.prank(owner);
+        token.transfer(user, amount);
+
+        vm.startPrank(user);
+        token.approve(address(staking), amount);
+        staking.stake(amount);
+
+        assertEq(staking.getStakedAmount(user), amount);
+        assertEq(token.balanceOf(newStaking), 0);
+
+        address newImplementation = address(new AgentStakingAccessStorage());
+
+        vm.startPrank(owner);
+        staking.upgradeToAndCall(newImplementation, "");
+
+        assertEq(staking.getStakedAmount(user), amount);
+        assertEq(token.balanceOf(newStaking), 0);
+
+        AgentStakingAccessStorage(address(staking)).migrate(newStaking, user);
+
+        assertEq(staking.getStakedAmount(user), 0);
+        assertEq(token.balanceOf(newStaking), amount);
+    }
+
     function _unstakeWarpAndClaim(address staker, uint256 amount) internal {
         vm.startPrank(staker);
 
@@ -466,7 +568,40 @@ contract AgentStakingTest is Test {
 }
 
 contract AgentStakingV2Mock is AgentStaking {
-    function test() external pure returns(bool) {
+    function test() public pure returns (bool) {
         return true;
     }
+}
+
+contract AgentStakingUnlock is AgentStaking {
+    function unlock_time() public view override returns (uint256) {
+        return 2 days;
+    }
+}
+
+contract AgentStakingDisabled is AgentStaking {
+    function stake(uint256 amount) public override {
+        revert("Staking is disabled");
+    }
+}
+
+contract AgentUnstaking is AgentStaking {
+    function unstake(uint256 amount) public override {
+        revert("Unstaking is disabled");
+    }
+}
+
+contract AgentStakingAccessStorage is AgentStaking {
+    using SafeERC20 for IERC20;
+
+    function migrate(address newStakingAddress, address user) public {
+        uint256 amount = stakes[user];
+
+        if (amount == 0) {
+            revert ("No stakes to migrate");
+        }
+
+        stakes[user] = 0;
+        agentToken.transfer(newStakingAddress, amount);
+    }
 }