Merge pull request #249 from Layr-Labs/test/strategy-manager-add-remo…

…ve-shares Additional StrategyManager test coverage
Layr-Labs · Oct 16, 2023 · e7bf6b2 · e7bf6b2
2 parents db916ae + c2eb3ff
commit e7bf6b2
Show file tree

Hide file tree

Showing 2 changed files with 277 additions and 273 deletions.
diff --git a/src/test/mocks/DelegationManagerMock.sol b/src/test/mocks/DelegationManagerMock.sol
@@ -3,6 +3,7 @@ pragma solidity =0.8.12;
 
 import "forge-std/Test.sol";
 import "../../contracts/interfaces/IDelegationManager.sol";
+import "../../contracts/interfaces/IStrategyManager.sol";
 
 
 contract DelegationManagerMock is IDelegationManager, Test {
@@ -135,4 +136,33 @@ contract DelegationManagerMock is IDelegationManager, Test {
         uint256[] calldata middlewareTimesIndexes,
         bool[] calldata receiveAsTokens
     ) external {}
+
+    // onlyDelegationManager functions in StrategyManager
+    function addShares(
+        IStrategyManager strategyManager,
+        address staker,
+        IStrategy strategy,
+        uint256 shares
+    ) external {
+        strategyManager.addShares(staker, strategy, shares);
+    }
+
+    function removeShares(
+        IStrategyManager strategyManager,
+        address staker,
+        IStrategy strategy,
+        uint256 shares
+    ) external {
+        strategyManager.removeShares(staker, strategy, shares);
+    }
+
+    function withdrawSharesAsTokens(
+        IStrategyManager strategyManager,
+        address recipient,
+        IStrategy strategy,
+        uint256 shares,
+        IERC20 token
+    ) external {
+        strategyManager.withdrawSharesAsTokens(recipient, strategy, shares, token);
+    }
 }