fix: missing tests for setMigrator

Author: horsefacts

test/IdRegistry/IdRegistry.migration.t.sol

@@ -32,6 +32,7 @@ contract IdRegistryTest is IdRegistryTestSuite {
     event AdminReset(uint256 indexed fid);
     event SetIdCounter(uint256 oldCounter, uint256 newCounter);
     event FreezeIdGateway(address idGateway);
+    event SetMigrator(address oldMigrator, address newMigrator);
 
     /*//////////////////////////////////////////////////////////////
                               PARAMETERS
@@ -53,6 +54,46 @@ contract IdRegistryTest is IdRegistryTestSuite {
         assertEq(idRegistry.isMigrated(), false);
     }
 
+    /*//////////////////////////////////////////////////////////////
+                             SET MIGRATOR
+    //////////////////////////////////////////////////////////////*/
+
+    function testFuzzOwnerCanSetMigrator(address migrator) public {
+        address oldMigrator = idRegistry.migrator();
+
+        vm.expectEmit();
+        emit SetMigrator(oldMigrator, migrator);
+        vm.prank(owner);
+        idRegistry.setMigrator(migrator);
+
+        assertEq(idRegistry.migrator(), migrator);
+    }
+
+    function testFuzzSetMigratorRevertsWhenMigrated(address migrator) public {
+        address oldMigrator = idRegistry.migrator();
+
+        vm.prank(oldMigrator);
+        idRegistry.migrate();
+
+        vm.prank(owner);
+        vm.expectRevert(IMigration.AlreadyMigrated.selector);
+        idRegistry.setMigrator(migrator);
+
+        assertEq(idRegistry.migrator(), oldMigrator);
+    }
+
+    function testFuzzSetMigratorRevertsWhenUnpaused(address migrator) public {
+        address oldMigrator = idRegistry.migrator();
+
+        vm.startPrank(owner);
+        idRegistry.unpause();
+        vm.expectRevert("Pausable: not paused");
+        idRegistry.setMigrator(migrator);
+        vm.stopPrank();
+
+        assertEq(idRegistry.migrator(), oldMigrator);
+    }
+
     /*//////////////////////////////////////////////////////////////
                                 MIGRATION
     //////////////////////////////////////////////////////////////*/

test/KeyRegistry/KeyRegistry.migration.t.sol

@@ -33,6 +33,7 @@ contract KeyRegistryTest is KeyRegistryTestSuite {
     event Remove(uint256 indexed fid, bytes indexed key, bytes keyBytes);
     event AdminReset(uint256 indexed fid, bytes indexed key, bytes keyBytes);
     event Migrated(uint256 indexed migratedAt);
+    event SetMigrator(address oldMigrator, address newMigrator);
 
     function testInitialGracePeriod() public {
         assertEq(keyRegistry.gracePeriod(), 1 days);
@@ -50,6 +51,46 @@ contract KeyRegistryTest is KeyRegistryTestSuite {
         assertEq(keyRegistry.isMigrated(), false);
     }
 
+    /*//////////////////////////////////////////////////////////////
+                             SET MIGRATOR
+    //////////////////////////////////////////////////////////////*/
+
+    function testFuzzOwnerCanSetMigrator(address migrator) public {
+        address oldMigrator = keyRegistry.migrator();
+
+        vm.expectEmit();
+        emit SetMigrator(oldMigrator, migrator);
+        vm.prank(owner);
+        keyRegistry.setMigrator(migrator);
+
+        assertEq(keyRegistry.migrator(), migrator);
+    }
+
+    function testFuzzSetMigratorRevertsWhenMigrated(address migrator) public {
+        address oldMigrator = keyRegistry.migrator();
+
+        vm.prank(oldMigrator);
+        keyRegistry.migrate();
+
+        vm.prank(owner);
+        vm.expectRevert(IMigration.AlreadyMigrated.selector);
+        keyRegistry.setMigrator(migrator);
+
+        assertEq(keyRegistry.migrator(), oldMigrator);
+    }
+
+    function testFuzzSetMigratorRevertsWhenUnpaused(address migrator) public {
+        address oldMigrator = keyRegistry.migrator();
+
+        vm.startPrank(owner);
+        keyRegistry.unpause();
+        vm.expectRevert("Pausable: not paused");
+        keyRegistry.setMigrator(migrator);
+        vm.stopPrank();
+
+        assertEq(keyRegistry.migrator(), oldMigrator);
+    }
+
     /*//////////////////////////////////////////////////////////////
                                 MIGRATION
     //////////////////////////////////////////////////////////////*/