add withdraw wait period

contracts/HydraStaking/HydraStaking.sol

@@ -204,7 +204,7 @@ contract HydraStaking is
     /**
      * @inheritdoc PenalizeableStaking
      */
-    function _afterWithdrawBannedFundsHook(address staker, uint256 withdrawnAmount) internal virtual override {
+    function _afterInitiatePenalizedFundsWithdrawal(address staker, uint256 withdrawnAmount) internal virtual override {
         _collectTokens(staker, withdrawnAmount);
     }
 

contracts/HydraStaking/modules/PenalizeableStaking/IPenalizeableStaking.sol

@@ -31,7 +31,8 @@ interface IPenalizeableStaking is IStaking {
     function penalizeStaker(address staker, PenalizedStakeDistribution[] calldata stakeDistributions) external;
 
     /**
-     * @notice Withdraws the funds of a banned validator
+     * @notice Register withdrawal of the penalized funds
+     * @dev The funds will be available for withdrawal after the withdrawal waiting period
      */
-    function withdrawBannedFunds() external;
+    function initiatePenalizedFundsWithdrawal() external;
 }

contracts/HydraStaking/modules/PenalizeableStaking/PenalizeableStaking.sol

@@ -51,16 +51,16 @@ abstract contract PenalizeableStaking is IPenalizeableStaking, HydraChainConnect
     /**
      * @inheritdoc IPenalizeableStaking
      */
-    function withdrawBannedFunds() external {
+    function initiatePenalizedFundsWithdrawal() external {
         uint256 leftToWithdraw = leftToWithdrawPerStaker[msg.sender];
         if (leftToWithdraw == 0) {
             revert NoFundsToWithdraw();
         }
 
         delete leftToWithdrawPerStaker[msg.sender];
 
-        _withdraw(msg.sender, leftToWithdraw);
-        _afterWithdrawBannedFundsHook(msg.sender, leftToWithdraw);
+        _registerWithdrawal(msg.sender, leftToWithdraw);
+        _afterInitiatePenalizedFundsWithdrawal(msg.sender, leftToWithdraw);
     }
 
     // _______________ Internal functions _______________
@@ -118,7 +118,7 @@ abstract contract PenalizeableStaking is IPenalizeableStaking, HydraChainConnect
      * @param withdrawnAmount The amount to withdraw
      */
     // solhint-disable-next-line no-empty-blocks
-    function _afterWithdrawBannedFundsHook(address staker, uint256 withdrawnAmount) internal virtual {}
+    function _afterInitiatePenalizedFundsWithdrawal(address staker, uint256 withdrawnAmount) internal virtual {}
 
     /**
      * @notice Handles the withdrawal of the given amount for the given account

docs/HydraStaking/HydraStaking.md

@@ -511,6 +511,17 @@ function initialize(StakerInit[] initialStakers, uint256 newMinStake, address go
 | rewardWalletAddr | address | undefined |
 | liquidToken | address | undefined |
 
+### initiatePenalizedFundsWithdrawal
+
+```solidity
+function initiatePenalizedFundsWithdrawal() external nonpayable
+```
+
+Register withdrawal of the penalized funds
+
+*The funds will be available for withdrawal after the withdrawal waiting period*
+
+
 ### lastDistribution
 
 ```solidity
@@ -1089,17 +1100,6 @@ Withdraws sender's withdrawable amount to specified address.
 |---|---|---|
 | to | address | Address to withdraw to |
 
-### withdrawBannedFunds
-
-```solidity
-function withdrawBannedFunds() external nonpayable
-```
-
-Withdraws the funds of a banned validator
-
-
-
-
 ### withdrawWaitPeriod
 
 ```solidity

docs/HydraStaking/IHydraStaking.md

@@ -200,6 +200,17 @@ Returns historical records of the staking rewards of the user
 |---|---|---|
 | _0 | StakingRewardsHistory[] | stakingRewardsHistory array with the historical records of the staking rewards of the user |
 
+### initiatePenalizedFundsWithdrawal
+
+```solidity
+function initiatePenalizedFundsWithdrawal() external nonpayable
+```
+
+Register withdrawal of the penalized funds
+
+*The funds will be available for withdrawal after the withdrawal waiting period*
+
+
 ### liquidToken
 
 ```solidity
@@ -478,17 +489,6 @@ Withdraws sender's withdrawable amount to specified address.
 |---|---|---|
 | to | address | Address to withdraw to |
 
-### withdrawBannedFunds
-
-```solidity
-function withdrawBannedFunds() external nonpayable
-```
-
-Withdraws the funds of a banned validator
-
-
-
-
 ### withdrawable
 
 ```solidity

docs/HydraStaking/modules/PenalizeableStaking/IPenalizeableStaking.md

@@ -53,6 +53,17 @@ Claims staking rewards for the sender.
 
 
 
+### initiatePenalizedFundsWithdrawal
+
+```solidity
+function initiatePenalizedFundsWithdrawal() external nonpayable
+```
+
+Register withdrawal of the penalized funds
+
+*The funds will be available for withdrawal after the withdrawal waiting period*
+
+
 ### penalizeStaker
 
 ```solidity
@@ -179,17 +190,6 @@ Withdraws sender's withdrawable amount to specified address.
 |---|---|---|
 | to | address | Address to withdraw to |
 
-### withdrawBannedFunds
-
-```solidity
-function withdrawBannedFunds() external nonpayable
-```
-
-Withdraws the funds of a banned validator
-
-
-
-
 ### withdrawable
 
 ```solidity

docs/HydraStaking/modules/PenalizeableStaking/PenalizeableStaking.md

@@ -183,6 +183,17 @@ function hydraChainContract() external view returns (contract IHydraChain)
 |---|---|---|
 | _0 | contract IHydraChain | undefined |
 
+### initiatePenalizedFundsWithdrawal
+
+```solidity
+function initiatePenalizedFundsWithdrawal() external nonpayable
+```
+
+Register withdrawal of the penalized funds
+
+*The funds will be available for withdrawal after the withdrawal waiting period*
+
+
 ### leftToWithdrawPerStaker
 
 ```solidity
@@ -483,17 +494,6 @@ Withdraws sender's withdrawable amount to specified address.
 |---|---|---|
 | to | address | Address to withdraw to |
 
-### withdrawBannedFunds
-
-```solidity
-function withdrawBannedFunds() external nonpayable
-```
-
-Withdraws the funds of a banned validator
-
-
-
-
 ### withdrawWaitPeriod
 
 ```solidity

test/HydraChain/Inspector.test.ts

@@ -1,7 +1,7 @@
 /* eslint-disable node/no-extraneous-import */
 import { ethers } from "hardhat";
 import { expect } from "chai";
-import { loadFixture, time, setBalance } from "@nomicfoundation/hardhat-network-helpers";
+import { loadFixture, time } from "@nomicfoundation/hardhat-network-helpers";
 
 import { calculatePenaltyByWeeks, commitEpochs } from "../helper";
 import { BAN_THRESHOLD, ERRORS, VESTING_DURATION_WEEKS, WEEK } from "../constants";
@@ -241,7 +241,10 @@ export function RunInspectorTests(): void {
       const liquidTokensThatMustBeCollected = (await hydraStaking.liquidityDebts(staker.address)).add(
         withdrawableAmount
       );
-      await expect(hydraStaking.connect(staker).withdrawBannedFunds(), "liquidTokens collect").to.changeTokenBalance(
+      await expect(
+        hydraStaking.connect(staker).initiatePenalizedFundsWithdrawal(),
+        "liquidTokens collect"
+      ).to.changeTokenBalance(
         LiquidityToken__factory.connect(await hydraStaking.liquidToken(), ethers.provider),
         staker.address,
         liquidTokensThatMustBeCollected.mul(-1)
@@ -352,7 +355,7 @@ export function RunInspectorTests(): void {
         withdrawableAmount
       );
       await expect(
-        hydraStaking.connect(inBanProcessValidator).withdrawBannedFunds(),
+        hydraStaking.connect(inBanProcessValidator).initiatePenalizedFundsWithdrawal(),
         "liquidTokens collect"
       ).to.changeTokenBalance(
         LiquidityToken__factory.connect(await hydraStaking.liquidToken(), ethers.provider),
@@ -606,54 +609,66 @@ export function RunInspectorTests(): void {
   });
 
   describe("Withdraw banned funds", function () {
-    it("should fail the withdrawal when there are no funds in the hydraStaking", async function () {
-      const { bannedValidator, hydraStaking } = await loadFixture(this.fixtures.bannedValidatorFixture);
-
-      // clear the contract's balance in order to force withdrawal fail
-      setBalance(hydraStaking.address, 0);
-      await expect(hydraStaking.connect(bannedValidator).withdrawBannedFunds()).to.be.revertedWithCustomError(
-        hydraStaking,
-        "WithdrawalFailed"
-      );
-    });
-
     it("should revert when trying to withdraw banned funds when there is none to withdraw", async function () {
       const { bannedValidator, hydraStaking } = await loadFixture(this.fixtures.bannedValidatorFixture);
 
-      await hydraStaking.connect(bannedValidator).withdrawBannedFunds();
+      await hydraStaking.connect(bannedValidator).initiatePenalizedFundsWithdrawal();
       expect(await hydraStaking.connect(bannedValidator).leftToWithdrawPerStaker(bannedValidator.address)).to.be.equal(
         0
       );
-      await expect(hydraStaking.connect(bannedValidator).withdrawBannedFunds()).to.be.revertedWithCustomError(
-        hydraStaking,
-        "NoFundsToWithdraw"
+      await expect(
+        hydraStaking.connect(bannedValidator).initiatePenalizedFundsWithdrawal()
+      ).to.be.revertedWithCustomError(hydraStaking, "NoFundsToWithdraw");
+    });
+
+    it("should initiatePenalizedFundsWithdrawal and clean up leftToWithdrawPerStaker & handle liquid tokens", async function () {
+      const { liquidToken, bannedValidator, stakedAmount, hydraStaking } = await loadFixture(
+        this.fixtures.bannedValidatorFixture
       );
+
+      const initiateWithdrawTx = await hydraStaking.connect(bannedValidator).initiatePenalizedFundsWithdrawal();
+
+      await expect(initiateWithdrawTx, "emit Transfer")
+        .to.emit(liquidToken, "Transfer")
+        .withArgs(bannedValidator.address, ethers.constants.AddressZero, stakedAmount);
+      expect(await liquidToken.balanceOf(bannedValidator.address), "lydra balanceOf").to.be.equal(0);
+
+      const withdrawalBalance = await hydraStaking.leftToWithdrawPerStaker(bannedValidator.address);
+      expect(withdrawalBalance, "withdrawalBalance.withdrawableAmount").to.be.equal(0);
+      expect(await hydraStaking.liquidityDebts(bannedValidator.address)).to.be.equal(0);
+    });
+
+    it("should fail to withdraw the banned funds if waiting withdraw period is not passed", async function () {
+      const { bannedValidator, hydraStaking } = await loadFixture(this.fixtures.bannedValidatorFixture);
+
+      await hydraStaking.connect(bannedValidator).initiatePenalizedFundsWithdrawal();
+
+      await expect(
+        hydraStaking.connect(bannedValidator).withdraw(bannedValidator.address)
+      ).to.be.revertedWithCustomError(hydraStaking, "NoWithdrawalAvailable");
     });
 
     it("should successfully withdraw the funds", async function () {
-      const { hydraChain, liquidToken, bannedValidator, stakedAmount, hydraStaking } = await loadFixture(
+      const { hydraChain, bannedValidator, stakedAmount, hydraStaking } = await loadFixture(
         this.fixtures.bannedValidatorFixture
       );
 
       const validatorBanPenalty = await hydraChain.validatorPenalty();
       const withdrawAmount = stakedAmount.sub(validatorBanPenalty);
-      const withdrawTx = await hydraStaking.connect(bannedValidator).withdrawBannedFunds();
+      await hydraStaking.connect(bannedValidator).initiatePenalizedFundsWithdrawal();
+
+      // increase time to pass the waiting period
+      await time.increase(WEEK * 2);
+
+      const withdrawTx = await hydraStaking.connect(bannedValidator).withdraw(bannedValidator.address);
 
-      await expect(withdrawTx, "emit Transfer")
-        .to.emit(liquidToken, "Transfer")
-        .withArgs(bannedValidator.address, ethers.constants.AddressZero, stakedAmount);
-      expect(await liquidToken.balanceOf(bannedValidator.address), "lydra balanceOf").to.be.equal(0);
       await expect(withdrawTx, "emit WithdrawalFinished")
         .to.emit(hydraStaking, "WithdrawalFinished")
         .withArgs(hydraStaking.address, bannedValidator.address, withdrawAmount);
       await expect(withdrawTx, "change balances").to.changeEtherBalances(
         [hydraStaking.address, bannedValidator.address],
         [withdrawAmount.mul(-1), withdrawAmount]
       );
-
-      const withdrawalBalance = await hydraStaking.leftToWithdrawPerStaker(bannedValidator.address);
-      expect(withdrawalBalance, "withdrawalBalance.withdrawableAmount").to.be.equal(0);
-      expect(await hydraStaking.liquidityDebts(bannedValidator.address)).to.be.equal(0);
     });
   });
 }