test: inherit staker in operator contract

Layr-Labs · Nov 16, 2023 · b8e423c · b8e423c
1 parent c4807d1
commit b8e423c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 23 deletions.
diff --git a/src/test/integration/Operator.sol b/src/test/integration/Operator.sol
@@ -1,14 +1,10 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity =0.8.12;
 
-import "forge-std/Test.sol";
-
+import "src/test/integration/Staker.sol";
 import "src/test/integration/GlobalRefs.sol";
 
-contract Operator is Test {
-    // Pointer to global reference contract
-    GlobalRefs globalRefs;
-
+contract Operator is Staker {
     // Array of addresses who are delegated to staker
     address[] public delegatedStaker;
     mapping(address => bool) public isDelegatedStaker;
@@ -19,9 +15,7 @@ contract Operator is Test {
     // Delegation Signer Private Key
     uint256 delegationSignerPrivateKey = uint256(0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80);
 
-    constructor(GlobalRefs _globalRefs) {
-        globalRefs = _globalRefs;
-    }
+    constructor(GlobalRefs _globalRefs) Staker (_globalRefs){}
 
     // Registration Functions
     function register() public {
@@ -58,6 +52,11 @@ contract Operator is Test {
         assertEq(globalRefs.delegationManager().delegatedTo(operator), operator, "operator not delegated to self");       
     }
 
+    // Override staker functions that operator shouldn't call
+    function delegate(address operator) public override {
+        revert("Operator cannot delegate since it's delegated to itself");
+    }
+
     // Helper functions to update stakers that are delegated to this operator
     function addDelegatedStaker(address staker) public {
         if (!isDelegatedStaker[staker]) {

diff --git a/src/test/integration/Scenario.t.sol b/src/test/integration/Scenario.t.sol
@@ -28,7 +28,8 @@ contract TestScenario is IntegrationTestRunner {
         operator.register();
 
         // Delegate To Operator
-        staker.delegate(operator);
+        staker.delegate(address(operator));
+        operator.addDelegatedStaker(address(staker));
 
         // Deposit into strategy
         staker.depositIntoStrategy(strategy1, strategy1Token, 1000);

diff --git a/src/test/integration/Staker.sol b/src/test/integration/Staker.sol
@@ -6,11 +6,10 @@ import "forge-std/Test.sol";
 import "src/contracts/interfaces/IDelegationManager.sol";
 
 import "src/test/integration/GlobalRefs.sol";
-import "src/test/integration/Operator.sol";
 
 contract Staker is Test {
     // Pointer to global reference contract
-    GlobalRefs globalRefs;
+    GlobalRefs public globalRefs;
 
     // Local Storage
     IDelegationManager.Withdrawal[] public queuedWithdrawals;
@@ -27,40 +26,37 @@ contract Staker is Test {
         globalRefs = _globalRefs;
     }
 
-    function delegate(Operator operator) public {
+    function delegate(address operator) public virtual {
         ISignatureUtils.SignatureWithExpiry memory approverSignatureAndExpiry;
         bytes32 emptySalt;
         _delegate(operator, approverSignatureAndExpiry, emptySalt);
 
         // Assert that salt is not spent
         assertFalse(globalRefs.delegationManager().delegationApproverSaltIsSpent(
-            globalRefs.delegationManager().delegationApprover(address(operator)),
+            globalRefs.delegationManager().delegationApprover(operator),
             emptySalt
         ), "salt somehow spent too early");
     }
 
-    function _delegate(Operator operator, ISignatureUtils.SignatureWithExpiry memory approverSignatureAndExpiry, bytes32 approverSalt) internal {
+    function _delegate(address operator, ISignatureUtils.SignatureWithExpiry memory approverSignatureAndExpiry, bytes32 approverSalt) internal {
         // Save pre-delegation state of operator 
         (IStrategy[] memory stakerStrategies, uint256[] memory stakerShares) = globalRefs.delegationManager().getDelegatableShares(staker);
         uint256[] memory operatorSharesBefore = new uint256[](stakerShares.length);
         for (uint256 i = 0; i < stakerShares.length; i++) {
-            operatorSharesBefore[i] = globalRefs.delegationManager().operatorShares(address(operator), stakerStrategies[i]);
+            operatorSharesBefore[i] = globalRefs.delegationManager().operatorShares(operator, stakerStrategies[i]);
         }
 
         // Delegate to operator
-        globalRefs.delegationManager().delegateTo(address(operator), approverSignatureAndExpiry, approverSalt);
+        globalRefs.delegationManager().delegateTo(operator, approverSignatureAndExpiry, approverSalt);
 
         // Checks
-        assertEq(globalRefs.delegationManager().delegatedTo(staker), address(operator), "staker delegated to the wrong address");
+        assertEq(globalRefs.delegationManager().delegatedTo(staker), operator, "staker delegated to the wrong address");
         for(uint256 i = 0; i < stakerShares.length; i++) {
-            assertEq(globalRefs.delegationManager().operatorShares(address(operator), stakerStrategies[i]), operatorSharesBefore[i] + stakerShares[i], "operator shares not updated correctly");
+            assertEq(globalRefs.delegationManager().operatorShares(operator, stakerStrategies[i]), operatorSharesBefore[i] + stakerShares[i], "operator shares not updated correctly");
         }
-
-        // Update operator's delegatedStaker array
-        operator.addDelegatedStaker(staker);
     }
 
-    function depositIntoStrategy(IStrategy strategy, IERC20 token, uint256 amount) public returns (uint256 shares){
+    function depositIntoStrategy(IStrategy strategy, IERC20 token, uint256 amount) public returns (uint256 shares) {
         // Save pre-deposit state
         uint256 stakerSharesBefore = _getStakerSharesForStrategy(strategy);
         uint256 operatorSharesBefore = _getOperatorSharesForStrategy(strategy);