refactor(StakeManager): add custom `StakeManager__FundsLocked` error (#29)

This commit is contained in:
r4bbit 2023-10-10 15:44:16 +02:00 committed by GitHub
parent eeffcfe7d7
commit 03d2dcf3e8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 34 additions and 3 deletions

View File

@ -7,5 +7,6 @@ MigrateTest:test_RevertWhen_SenderIsNotVault() (gas: 10629)
StakeManagerTest:testDeployment() (gas: 26172) StakeManagerTest:testDeployment() (gas: 26172)
StakeTest:testDeployment() (gas: 26172) StakeTest:testDeployment() (gas: 26172)
StakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10638) StakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10638)
UnstakeTest:testDeployment() (gas: 26172) UnstakeTest:testDeployment() (gas: 26355)
UnstakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10575) UnstakeTest:test_RevertWhen_FundsLocked() (gas: 973435)
UnstakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10609)

View File

@ -8,6 +8,7 @@ import { StakeVault } from "./StakeVault.sol";
contract StakeManager is Ownable { contract StakeManager is Ownable {
error StakeManager__SenderIsNotVault(); error StakeManager__SenderIsNotVault();
error StakeManager__FundsLocked();
struct Account { struct Account {
uint256 lockUntil; uint256 lockUntil;
@ -74,7 +75,9 @@ contract StakeManager is Ownable {
*/ */
function unstake(uint256 _amount) external onlyVault { function unstake(uint256 _amount) external onlyVault {
Account storage account = accounts[msg.sender]; Account storage account = accounts[msg.sender];
require(account.lockUntil <= block.timestamp, "Funds are locked"); if (account.lockUntil > block.timestamp) {
revert StakeManager__FundsLocked();
}
processAccount(account, currentEpoch); processAccount(account, currentEpoch);
uint256 reducedMultiplier = (_amount * account.multiplier) / account.balance; uint256 reducedMultiplier = (_amount * account.multiplier) / account.balance;
account.multiplier -= reducedMultiplier; account.multiplier -= reducedMultiplier;

View File

@ -1,10 +1,13 @@
// SPDX-License-Identifier: UNLICENSED // SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.19; pragma solidity ^0.8.19;
import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import { Test } from "forge-std/Test.sol"; import { Test } from "forge-std/Test.sol";
import { Deploy } from "../script/Deploy.s.sol"; import { Deploy } from "../script/Deploy.s.sol";
import { DeploymentConfig } from "../script/DeploymentConfig.s.sol"; import { DeploymentConfig } from "../script/DeploymentConfig.s.sol";
import { StakeManager } from "../contracts/StakeManager.sol"; import { StakeManager } from "../contracts/StakeManager.sol";
import { StakeVault } from "../contracts/StakeVault.sol";
contract StakeManagerTest is Test { contract StakeManagerTest is Test {
DeploymentConfig internal deploymentConfig; DeploymentConfig internal deploymentConfig;
@ -12,6 +15,7 @@ contract StakeManagerTest is Test {
address internal stakeToken; address internal stakeToken;
address internal deployer; address internal deployer;
address internal testUser = makeAddr("testUser");
function setUp() public virtual { function setUp() public virtual {
Deploy deployment = new Deploy(); Deploy deployment = new Deploy();
@ -29,6 +33,14 @@ contract StakeManagerTest is Test {
assertEq(address(stakeManager.oldManager()), address(0)); assertEq(address(stakeManager.oldManager()), address(0));
assertEq(stakeManager.totalSupply(), 0); assertEq(stakeManager.totalSupply(), 0);
} }
function _createTestVault(address owner) internal returns (StakeVault vault) {
vm.prank(owner);
vault = new StakeVault(owner, ERC20(stakeToken), stakeManager);
vm.prank(deployer);
stakeManager.setVault(address(vault).codehash);
}
} }
contract StakeTest is StakeManagerTest { contract StakeTest is StakeManagerTest {
@ -51,6 +63,21 @@ contract UnstakeTest is StakeManagerTest {
vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector); vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
stakeManager.unstake(100); stakeManager.unstake(100);
} }
function test_RevertWhen_FundsLocked() public {
// ensure user has funds
deal(stakeToken, testUser, 1000);
StakeVault userVault = _createTestVault(testUser);
vm.startPrank(testUser);
ERC20(stakeToken).approve(address(userVault), 100);
uint256 lockTime = 1 days;
userVault.stake(100, lockTime);
vm.expectRevert(StakeManager.StakeManager__FundsLocked.selector);
userVault.unstake(100);
}
} }
contract LockTest is StakeManagerTest { contract LockTest is StakeManagerTest {