Make unmigration lock a boolean for simplicity and future maintenance

0xPolygon · Sep 1, 2023 · b9c3cfe · b9c3cfe
1 parent 3b9abc2
commit b9c3cfe
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 16 deletions.
diff --git a/src/PolygonMigration.sol b/src/PolygonMigration.sol
@@ -19,10 +19,10 @@ contract PolygonMigration is Ownable2Step, IPolygonMigration {
     IERC20 public polygon;
     IERC20 public immutable matic;
     uint256 public releaseTimestamp;
-    uint256 public unmigrationLock;
+    bool public unmigrationLocked;
 
     modifier ifUnmigrationUnlocked() {
-        if (unmigrationLock != 0) revert UnmigrationLocked();
+        if (unmigrationLocked) revert UnmigrationLocked();
         _;
     }
 
@@ -109,12 +109,10 @@ contract PolygonMigration is Ownable2Step, IPolygonMigration {
 
     /// @notice Allows governance to lock or unlock the unmigration process
     /// @dev The function does not do any validation since governance can update the unmigration process if required
-    /// @param unmigrationLock_ New unmigration lock status
-    function updateUnmigrationLock(
-        uint256 unmigrationLock_
-    ) external onlyOwner {
-        unmigrationLock = unmigrationLock_;
-        emit UnmigrationLockUpdated(unmigrationLock_);
+    /// @param unmigrationLocked_ New unmigration lock status
+    function updateUnmigrationLock(bool unmigrationLocked_) external onlyOwner {
+        unmigrationLocked = unmigrationLocked_;
+        emit UnmigrationLockUpdated(unmigrationLocked_);
     }
 
     /// @notice Allows governance to release the remaining POL tokens after the migration period has elapsed

diff --git a/src/interfaces/IPolygonMigration.sol b/src/interfaces/IPolygonMigration.sol
@@ -10,7 +10,7 @@ interface IPolygonMigration {
 
     event Migrated(address indexed account, uint256 amount);
     event Unmigrated(address indexed account, uint256 amount);
-    event UnmigrationLockUpdated(uint256 lock);
+    event UnmigrationLockUpdated(bool lock);
     event ReleaseTimestampUpdated(uint256 timestamp);
     event Released(uint256 polAmount, uint256 maticAmount);
 

diff --git a/test/PolygonMigration.t.sol b/test/PolygonMigration.t.sol
@@ -113,16 +113,12 @@ contract PolygonMigrationTest is Test {
         assertEq(matic.balanceOf(user), amount2);
     }
 
-    function testRevert_Unmigrate(
-        address user,
-        uint256 amount,
-        uint256 unmigrationLock
-    ) external {
+    function testRevert_Unmigrate(address user, uint256 amount) external {
+        bool unmigrationLock = true;
         vm.assume(
             amount <= 10000000000 * 10 ** 18 &&
                 user != address(0) &&
-                user != address(migration) &&
-                unmigrationLock != 0
+                user != address(migration)
         );
         matic.mint(user, amount);
         vm.startPrank(user);