test(KeyRegistry): add test coverage for Ownable2Step

script/LocalDeploy.s.sol

@@ -36,16 +36,11 @@ contract LocalDeploy is Script {
 
         vm.startBroadcast();
         (AggregatorV3Interface priceFeed, AggregatorV3Interface uptimeFeed) = _getOrDeployPriceFeeds();
-        IdRegistry idRegistry = new IdRegistry{salt: ID_REGISTRY_CREATE2_SALT}(
-            migrator,
-            initialIdRegistryOwner
+        IdRegistry idRegistry = new IdRegistry{salt: ID_REGISTRY_CREATE2_SALT}(migrator, initialIdRegistryOwner);
+        KeyRegistry keyRegistry = new KeyRegistry{salt: KEY_REGISTRY_CREATE2_SALT}(
+            address(idRegistry), migrator, initialKeyRegistryOwner, 1000
         );
-        KeyRegistry keyRegistry = new KeyRegistry{
-            salt: KEY_REGISTRY_CREATE2_SALT
-        }(address(idRegistry), migrator, initialKeyRegistryOwner, 1000);
-        StorageRegistry storageRegistry = new StorageRegistry{
-            salt: STORAGE_RENT_CREATE2_SALT
-        }(
+        StorageRegistry storageRegistry = new StorageRegistry{salt: STORAGE_RENT_CREATE2_SALT}(
             priceFeed,
             uptimeFeed,
             INITIAL_USD_UNIT_PRICE,

test/Bundler/Bundler.gas.t.sol

@@ -20,9 +20,7 @@ contract BundleRegistryGasUsageTest is BundlerTestSuite {
             bytes memory sig = _signRegister(i, account, address(0), type(uint40).max);
             uint256 price = bundler.price(1);
 
-            IBundler.SignerParams[] memory signers = new IBundler.SignerParams[](
-                0
-            );
+            IBundler.SignerParams[] memory signers = new IBundler.SignerParams[](0);
 
             vm.deal(account, 10_000 ether);
             vm.prank(account);

test/Bundler/Bundler.t.sol

@@ -41,9 +41,7 @@ contract BundlerTest is BundlerTestSuite {
         uint256 deadline,
         uint256 numSigners
     ) internal returns (IBundler.SignerParams[] memory) {
-        IBundler.SignerParams[] memory signers = new IBundler.SignerParams[](
-            numSigners
-        );
+        IBundler.SignerParams[] memory signers = new IBundler.SignerParams[](numSigners);
         uint256 nonce = keyRegistry.nonces(account);
 
         // The duplication below is ugly but necessary to work around a stack too deep error.

test/Bundler/BundlerTestSuite.sol

@@ -22,10 +22,7 @@ abstract contract BundlerTestSuite is IdGatewayTestSuite {
         keyRegistry.setKeyGateway(address(keyGateway));
 
         // Set up the BundleRegistry
-        bundler = new Bundler(
-            address(idGateway),
-            address(keyGateway)
-        );
+        bundler = new Bundler(address(idGateway), address(keyGateway));
 
         addKnownContract(address(keyGateway));
         addKnownContract(address(bundler));

test/IdGateway/IdGatewayTestSuite.sol

@@ -16,11 +16,7 @@ abstract contract IdGatewayTestSuite is StorageRegistryTestSuite, KeyRegistryTes
     function setUp() public virtual override(StorageRegistryTestSuite, KeyRegistryTestSuite) {
         super.setUp();
 
-        idGateway = new IdGateway(
-            address(idRegistry),
-            address(storageRegistry),
-            owner
-        );
+        idGateway = new IdGateway(address(idRegistry), address(storageRegistry), owner);
 
         vm.startPrank(owner);
         idRegistry.setIdGateway(address(idGateway));

test/KeyRegistry/KeyRegistry.t.sol

@@ -34,6 +34,7 @@ contract KeyRegistryTest is KeyRegistryTestSuite {
     event SetKeyGateway(address oldKeyGateway, address newKeyGateway);
     event SetMaxKeysPerFid(uint256 oldMax, uint256 newMax);
     event FreezeKeyGateway(address keyGateway);
+    event OwnershipTransferStarted(address indexed previousOwner, address indexed newOwner);
 
     function testInitialIdRegistry() public {
         assertEq(address(keyRegistry.idRegistry()), address(idRegistry));
@@ -43,6 +44,10 @@ contract KeyRegistryTest is KeyRegistryTestSuite {
         assertEq(keyRegistry.owner(), owner);
     }
 
+    function testInitialPendingOwner() public {
+        assertEq(keyRegistry.pendingOwner(), address(0));
+    }
+
     function testVersion() public {
         assertEq(keyRegistry.VERSION(), "2023.11.15");
     }
@@ -1104,6 +1109,48 @@ contract KeyRegistryTest is KeyRegistryTestSuite {
         keyRegistry.keysOf(0, IKeyRegistry.KeyState.NULL, 0, 1);
     }
 
+    /*//////////////////////////////////////////////////////////////
+                        TRANSFER OWNERSHIP 2 STEP
+    //////////////////////////////////////////////////////////////*/
+
+    function testTransferOwnershipRevertNonOwner(address newOwner) public {
+        address randomUser = makeAddr("randomUser");
+
+        vm.prank(randomUser);
+        vm.expectRevert("Ownable: caller is not the owner");
+        keyRegistry.transferOwnership(newOwner);
+
+        assertEq(keyRegistry.pendingOwner(), address(0));
+        assertEq(keyRegistry.owner(), owner);
+    }
+
+    function testTransferOwnershipRevertNonNewOwnerAcceptance(address newOwner) public {
+        address randomUser = makeAddr("randomUser");
+
+        vm.prank(owner);
+        keyRegistry.transferOwnership(newOwner);
+
+        vm.prank(randomUser);
+        vm.expectRevert("Ownable2Step: caller is not the new owner");
+        keyRegistry.acceptOwnership();
+        assertEq(keyRegistry.pendingOwner(), newOwner);
+        assertEq(keyRegistry.owner(), owner);
+    }
+
+    function testFuzzTransferOwnership(address newOwner) public {
+        vm.prank(owner);
+        vm.expectEmit(address(keyRegistry));
+        emit OwnershipTransferStarted(owner, newOwner);
+        keyRegistry.transferOwnership(newOwner);
+
+        assertEq(keyRegistry.pendingOwner(), newOwner);
+
+        vm.prank(newOwner);
+        keyRegistry.acceptOwnership();
+
+        assertEq(keyRegistry.owner(), newOwner);
+    }
+
     /*//////////////////////////////////////////////////////////////
                                  HELPERS
     //////////////////////////////////////////////////////////////*/

test/KeyRegistry/KeyRegistryTestHelpers.sol

@@ -42,9 +42,7 @@ library BulkAddDataBuilder {
         bytes memory metadata
     ) internal pure returns (KeyRegistry.BulkAddData[] memory) {
         KeyRegistry.BulkAddKey[] memory keys = addData[index].keys;
-        KeyRegistry.BulkAddKey[] memory newKeys = new KeyRegistry.BulkAddKey[](
-            keys.length + 1
-        );
+        KeyRegistry.BulkAddKey[] memory newKeys = new KeyRegistry.BulkAddKey[](keys.length + 1);
 
         for (uint256 i; i < keys.length; i++) {
             newKeys[i] = keys[i];
@@ -65,9 +63,7 @@ library BulkResetDataBuilder {
         KeyRegistry.BulkResetData[] memory resetData,
         uint256 fid
     ) internal pure returns (KeyRegistry.BulkResetData[] memory) {
-        KeyRegistry.BulkResetData[] memory newData = new KeyRegistry.BulkResetData[](
-                resetData.length + 1
-            );
+        KeyRegistry.BulkResetData[] memory newData = new KeyRegistry.BulkResetData[](resetData.length + 1);
         for (uint256 i; i < resetData.length; i++) {
             newData[i] = resetData[i];
         }