From c361effb835428dbdaa8324f22ec32742a5df6fa Mon Sep 17 00:00:00 2001
From: Michael Sun <michaelsun97@gmail.com>
Date: Sat, 8 Feb 2025 19:57:38 -0500
Subject: [PATCH] test: integration and unit test fixes

---
 src/SlashingRegistryCoordinator.sol    | 10 ++--------
 test/integration/OperatorSetUser.t.sol | 24 ++++++++++++++++++++++++
 test/mocks/AllocationManagerMock.sol   | 23 ++++++++++++++++++++++-
 3 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/src/SlashingRegistryCoordinator.sol b/src/SlashingRegistryCoordinator.sol
index 8292d4ee..1d3ec050 100644
--- 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
index a3a5ca0d..19ef0b42 100644
--- 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
index 22b92b74..23a14e4c 100644
--- 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;
+    }
+}