feat: address TODOs in tests

Layr-Labs · Jan 9, 2024 · 1f1a840 · 1f1a840
1 parent 278b660
commit 1f1a840
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 6 deletions.
diff --git a/test/mocks/StakeRegistryMock.sol b/test/mocks/StakeRegistryMock.sol
@@ -9,6 +9,12 @@ import "../../src/interfaces/IRegistryCoordinator.sol";
  * @author Layr Labs, Inc.
  */
 contract StakeRegistryMock is IStakeRegistry {
+    // bitmap returned by the mocked `updateOperatorStake` function
+    uint192 updateOperatorStakeReturnBitmap;
+
+    function set_updateOperatorStakeReturnBitmap(uint192 newValue) external {
+        updateOperatorStakeReturnBitmap = newValue;
+    }
 
     function registryCoordinator() external view returns (address) {}
 
@@ -195,10 +201,12 @@ contract StakeRegistryMock is IStakeRegistry {
      * added to the
      */
     function updateOperatorStake(
-        address operator, 
-        bytes32 operatorId, 
-        bytes calldata quorumNumbers
-    ) external returns (uint192) {}
+        address /*operator*/, 
+        bytes32 /*operatorId*/, 
+        bytes calldata /*quorumNumbers*/
+    ) external returns (uint192) {
+        return updateOperatorStakeReturnBitmap;
+    }
 
     function getMockOperatorId(address operator) external pure returns(bytes32) {
         return bytes32(uint256(keccak256(abi.encodePacked(operator, "operatorId"))));

diff --git a/test/unit/RegistryCoordinatorUnit.t.sol b/test/unit/RegistryCoordinatorUnit.t.sol
@@ -1152,7 +1152,6 @@ contract RegistryCoordinatorUnitTests_DeregisterOperator_EjectOperator is Regist
         registryCoordinator.ejectOperator(defaultOperator, quorumNumbers);
     }
 
-    // TODO: this test currently fails. either need to document behavior + modify the test, or modify the code
     function test_getQuorumBitmapIndicesAtBlockNumber_revert_notRegistered() public {
         uint32 blockNumber;
         bytes32[] memory operatorIds = new bytes32[](1);
@@ -1522,10 +1521,47 @@ contract RegistryCoordinatorUnitTests_UpdateOperators is RegistryCoordinatorUnit
         address[] memory operatorsToUpdate = new address[](1);
         operatorsToUpdate[0] = defaultOperator;
 
-        // TODO: any additional checks here? mostly this just calls the StakeRegistry, so more appropriate for an integration test
         registryCoordinator.updateOperators(operatorsToUpdate);
     }
 
+    // @notice tests the `updateOperators` function with a single registered operator as input
+    // @dev also sets up return data from the StakeRegistry
+    function testFuzz_updateOperators_singleOperator(uint192 registrationBitmap, uint192 mockReturnData) public {
+        // filter fuzzed inputs to only valid inputs
+        cheats.assume(registrationBitmap != 0);
+        mockReturnData = (mockReturnData & registrationBitmap);
+        emit log_named_uint("mockReturnData", mockReturnData);
+
+        // register the default operator
+        ISignatureUtils.SignatureWithSaltAndExpiry memory emptySig;
+        uint32 registrationBlockNumber = 100;
+        bytes memory quorumNumbers = BitmapUtils.bitmapToBytesArray(registrationBitmap);
+        for (uint256 i = 0; i < quorumNumbers.length; ++i) {
+            stakeRegistry.setOperatorWeight(uint8(quorumNumbers[i]), defaultOperator, defaultStake);            
+        }
+        cheats.startPrank(defaultOperator);
+        cheats.roll(registrationBlockNumber);
+        registryCoordinator.registerOperator(quorumNumbers, defaultSocket, pubkeyRegistrationParams, emptySig);
+
+        address[] memory operatorsToUpdate = new address[](1);
+        operatorsToUpdate[0] = defaultOperator;
+
+        uint192 quorumBitmapBefore = registryCoordinator.getCurrentQuorumBitmap(defaultOperatorId);
+        assertEq(quorumBitmapBefore, registrationBitmap, "operator bitmap somehow incorrect");
+
+        // make the stake registry return info that the operator should be removed from quorums
+        uint192 quorumBitmapToRemove = mockReturnData;
+        bytes memory quorumNumbersToRemove = BitmapUtils.bitmapToBytesArray(quorumBitmapToRemove);
+        for (uint256 i = 0; i < quorumNumbersToRemove.length; ++i) {
+            stakeRegistry.setOperatorWeight(uint8(quorumNumbersToRemove[i]), defaultOperator, 0);            
+        }
+        uint256 expectedQuorumBitmap = BitmapUtils.minus(quorumBitmapBefore, quorumBitmapToRemove);
+
+        registryCoordinator.updateOperators(operatorsToUpdate);
+        uint192 quorumBitmapAfter = registryCoordinator.getCurrentQuorumBitmap(defaultOperatorId);
+        assertEq(expectedQuorumBitmap, quorumBitmapAfter, "quorum bitmap did not update correctly");
+    }
+
     // @notice tests the `updateOperators` function with a single *un*registered operator as input
     function test_updateOperators_unregisteredOperator() public view {
         address[] memory operatorsToUpdate = new address[](1);