test: integration and unit test fixes

Layr-Labs · Feb 9, 2025 · c361eff · c361eff
1 parent eb7aa3d
commit c361eff
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 9 deletions.
diff --git a/src/SlashingRegistryCoordinator.sol b/src/SlashingRegistryCoordinator.sol
@@ -245,9 +245,7 @@ contract SlashingRegistryCoordinator is
             for (uint256 j = 0; j < quorumNumbers.length; j++) {
                 // update the operator's stake for each quorum
                 _updateStakesAndDeregisterLoiterers(
-                    singleOperator,
-                    singleOperatorId,
-                    uint8(quorumNumbers[j])
+                    singleOperator, singleOperatorId, uint8(quorumNumbers[j])
                 );
             }
         }
@@ -299,11 +297,7 @@ contract SlashingRegistryCoordinator is
                 prevOperatorAddress = operator;
             }
 
-            _updateStakesAndDeregisterLoiterers(
-                currQuorumOperators,
-                operatorIds,
-                quorumNumber
-            );
+            _updateStakesAndDeregisterLoiterers(currQuorumOperators, operatorIds, quorumNumber);
 
             // Update timestamp that all operators in quorum have been updated all at once
             quorumUpdateBlockNumber[quorumNumber] = block.number;

diff --git a/test/integration/OperatorSetUser.t.sol b/test/integration/OperatorSetUser.t.sol
@@ -117,6 +117,30 @@ contract OperatorSetUser is User {
         allocationManager.deregisterFromOperatorSets({params: deregisterParams});
     }
 
+    /// @dev Uses updateOperators to update this user's stake
+    function updateStakes() public virtual override createSnapshot {
+        _log("updateStakes (updateOperators)");
+
+        // get all quorums this operator is registered for
+        uint192 currentBitmap = slashingRegistryCoordinator.getCurrentQuorumBitmap(operatorId);
+        bytes memory quorumNumbers = currentBitmap.bitmapToBytesArray();
+
+        // get all operators in those quorums
+        address[][] memory operatorsPerQuorum = new address[][](quorumNumbers.length);
+        for (uint256 i = 0; i < quorumNumbers.length; i++) {
+            bytes32[] memory operatorIds = indexRegistry.getOperatorListAtBlockNumber(
+                uint8(quorumNumbers[i]), uint32(block.number)
+            );
+            operatorsPerQuorum[i] = new address[](operatorIds.length);
+            for (uint256 j = 0; j < operatorIds.length; j++) {
+                operatorsPerQuorum[i][j] = blsApkRegistry.pubkeyHashToOperator(operatorIds[j]);
+            }
+
+            operatorsPerQuorum[i] = Sort.sortAddresses(operatorsPerQuorum[i]);
+        }
+        slashingRegistryCoordinator.updateOperatorsForQuorum(operatorsPerQuorum, quorumNumbers);
+    }
+
     function _getOperatorSetIds(
         bytes memory quorums
     ) internal pure returns (uint32[] memory) {

diff --git a/test/mocks/AllocationManagerMock.sol b/test/mocks/AllocationManagerMock.sol
@@ -151,4 +151,25 @@ contract AllocationManagerIntermediate is IAllocationManager {
     ) external view virtual returns (bool) {}
 }
 
-contract AllocationManagerMock is AllocationManagerIntermediate {}
+contract AllocationManagerMock is AllocationManagerIntermediate {
+    bool defaultIsMemberOfOperatorSet;
+
+    constructor() {
+        // Return true by default to so that the SlashingRegistryCoordinator won't force deregister
+        // an operator for every quorum when a `updateOperators`,`updateOperatorsForQuorum` call is made
+        defaultIsMemberOfOperatorSet = true;
+    }
+
+    function isMemberOfOperatorSet(
+        address operator,
+        OperatorSet memory operatorSet
+    ) external view virtual override returns (bool) {
+        return defaultIsMemberOfOperatorSet;
+    }
+
+    function setDefaultIsMemberOfOperatorSet(
+        bool isMemberOfOperatorSet
+    ) external {
+        defaultIsMemberOfOperatorSet = isMemberOfOperatorSet;
+    }
+}