add properties for destroyStaker()

bool _allowFarmWithdraw: gives ability to withdraw vested tokens bool _allowClaimRewards: gives ability to withdraw accumulated rewards
bitcoin-portal · Jan 17, 2024 · 8b0f9dd · 8b0f9dd
1 parent a394c48
commit 8b0f9dd
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 8 deletions.
diff --git a/contracts/TimeLockFarmV2Dual.sol b/contracts/TimeLockFarmV2Dual.sol
@@ -474,23 +474,33 @@ contract TimeLockFarmV2Dual is TokenWrapper {
      * for the specified wallet address if leaving company or...
      */
     function destroyStaker(
+        bool _allowFarmWithdraw,
+        bool _allowClaimRewards,
         address _withdrawAddress
     )
         external
         onlyOwner
         updateFarm()
         updateAddy(_withdrawAddress)
     {
-        _claimReward(
-            _withdrawAddress
-        );
+        if (_allowFarmWithdraw == true) {
+            _farmWithdraw(
+                _withdrawAddress,
+                unlockable(
+                    _withdrawAddress
+                )
+            );
+        }
 
-        _farmWithdraw(
-            _withdrawAddress,
-            unlockable(
+        if (_allowClaimRewards == true) {
+            _claimReward(
                 _withdrawAddress
-            )
-        );
+            );
+        } else {
+            _takeRewards(
+                _withdrawAddress
+            );
+        }
 
         uint256 i;
         uint256 remainingStakes = stakes[_withdrawAddress].length;
@@ -506,6 +516,28 @@ contract TimeLockFarmV2Dual is TokenWrapper {
         );
     }
 
+    function _takeRewards(
+        address _fromAddress
+    )
+        internal
+    {
+        userRewardsA[ownerAddress] += userRewardsA[
+            _fromAddress
+        ];
+
+        userRewardsB[ownerAddress] += userRewardsB[
+            _fromAddress
+        ];
+
+        delete userRewardsA[
+            _fromAddress
+        ];
+
+        delete userRewardsB[
+            _fromAddress
+        ];
+    }
+
     /**
      * @dev Performs withdrawal of staked token from the farm
      */

diff --git a/contracts/TimeLockFarmV2Dual.t.sol b/contracts/TimeLockFarmV2Dual.t.sol
@@ -138,10 +138,14 @@ contract TimeLockFarmV2DualTest is Test {
         _simpleForwardTime();
 
         farm.destroyStaker(
+            true,
+            true,
             ADMIN_ADDRESS
         );
 
         farm.destroyStaker(
+            true,
+            true,
             0x6fEeB0c3E25E5dEf17BC7274406F0674B8237038
         );