From c48a185519606272253c41fcf41fad924f62edab Mon Sep 17 00:00:00 2001
From: QEDK <1994constant@gmail.com>
Date: Mon, 7 Aug 2023 00:00:20 +0400
Subject: [PATCH] test: add unmigration tests

---
 test/PolygonMigration.t.sol | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/test/PolygonMigration.t.sol b/test/PolygonMigration.t.sol
index a3f6b03..192e312 100644
--- a/test/PolygonMigration.t.sol
+++ b/test/PolygonMigration.t.sol
@@ -46,10 +46,29 @@ contract PolygonMigrationTest is Test {
         migration.migrate(amount);
 
         assertEq(matic.balanceOf(user), 0);
-        assertEq(matic.balanceOf(0x000000000000000000000000000000000000dEaD), amount);
+        assertEq(matic.balanceOf(address(migration)), amount);
         assertEq(polygon.balanceOf(user), amount);
     }
 
+    function test_Unmigrate(address user, uint256 amount, uint256 amount2) external {
+        vm.assume(amount <= 10000000000 * 10 ** 18 && amount2 <= amount && user != address(0));
+        matic.mint(user, amount);
+        vm.startPrank(user);
+        matic.approve(address(migration), amount);
+        migration.migrate(amount);
+
+        assertEq(matic.balanceOf(user), 0);
+        assertEq(matic.balanceOf(address(migration)), amount);
+        assertEq(polygon.balanceOf(user), amount);
+
+        polygon.approve(address(migration), amount2);
+        migration.unmigrate(amount2);
+
+        assertEq(polygon.balanceOf(user), amount - amount2);
+        assertEq(matic.balanceOf(address(migration)), amount - amount2);
+        assertEq(matic.balanceOf(user), amount2);
+    }
+
     function testRevert_UpdateReleaseTimestampOnlyGovernance(address user, uint256 timestamp) external {
         vm.assume(timestamp >= block.timestamp && user != governance);
         vm.expectRevert("Ownable: caller is not the owner");