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

Also introduce tests that ensure the error is actually emitted.

.gas-snapshot

@@ -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)

contracts/StakeManager.sol

@@ -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();
+        }
         _;
     }
 

test/StakeManager.t.sol

@@ -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();
+    }
+}