refactor(StakeManager): use custom error in `onlyVault` modifier (#28)

Also introduce tests that ensure the error is actually emitted.
This commit is contained in:
r4bbit 2023-10-10 15:32:46 +02:00 committed by GitHub
parent 725d380547
commit eeffcfe7d7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 71 additions and 2 deletions

View File

@ -1 +1,11 @@
StakeManagerTest:testDeployment() (gas: 26172)
LeaveTest:testDeployment() (gas: 26172)
LeaveTest:test_RevertWhen_SenderIsNotVault() (gas: 10562)
LockTest:testDeployment() (gas: 26172)
LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10573)
MigrateTest:testDeployment() (gas: 26172)
MigrateTest:test_RevertWhen_SenderIsNotVault() (gas: 10629)
StakeManagerTest:testDeployment() (gas: 26172)
StakeTest:testDeployment() (gas: 26172)
StakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10638)
UnstakeTest:testDeployment() (gas: 26172)
UnstakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10575)

View File

@ -7,6 +7,8 @@ import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
import { StakeVault } from "./StakeVault.sol";
contract StakeManager is Ownable {
error StakeManager__SenderIsNotVault();
struct Account {
uint256 lockUntil;
uint256 balance;
@ -40,7 +42,9 @@ contract StakeManager is Ownable {
ERC20 public immutable stakedToken;
modifier onlyVault() {
require(isVault[msg.sender.codehash], "Not a vault");
if (!isVault[msg.sender.codehash]) {
revert StakeManager__SenderIsNotVault();
}
_;
}

View File

@ -30,3 +30,58 @@ contract StakeManagerTest is Test {
assertEq(stakeManager.totalSupply(), 0);
}
}
contract StakeTest is StakeManagerTest {
function setUp() public override {
StakeManagerTest.setUp();
}
function test_RevertWhen_SenderIsNotVault() public {
vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
stakeManager.stake(100, 1);
}
}
contract UnstakeTest is StakeManagerTest {
function setUp() public override {
StakeManagerTest.setUp();
}
function test_RevertWhen_SenderIsNotVault() public {
vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
stakeManager.unstake(100);
}
}
contract LockTest is StakeManagerTest {
function setUp() public override {
StakeManagerTest.setUp();
}
function test_RevertWhen_SenderIsNotVault() public {
vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
stakeManager.lock(100);
}
}
contract LeaveTest is StakeManagerTest {
function setUp() public override {
StakeManagerTest.setUp();
}
function test_RevertWhen_SenderIsNotVault() public {
vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
stakeManager.leave();
}
}
contract MigrateTest is StakeManagerTest {
function setUp() public override {
StakeManagerTest.setUp();
}
function test_RevertWhen_SenderIsNotVault() public {
vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
stakeManager.migrate();
}
}