2372 stake table add tests to the withdrawfunds function

contracts/src/StakeTable.sol

@@ -152,7 +152,6 @@ contract StakeTable is AbstractStakeTable {
 
     /// @notice Get the next available epoch and queue size in that epoch
     /// TODO modify this according to the current spec
-    /// TODO modify this according to the current spec
     function nextRegistrationEpoch() external view override returns (uint64, uint64) {
         uint64 epoch;
         uint64 queueSize;
@@ -175,22 +174,19 @@ contract StakeTable is AbstractStakeTable {
     // @param queueSize current size of the registration queue (after insertion of new element in
     // the queue)
     /// TODO modify this according to the current spec
-    /// TODO modify this according to the current spec
     function appendRegistrationQueue(uint64 epoch, uint64 queueSize) private {
         firstAvailableRegistrationEpoch = epoch;
         _numPendingRegistrations = queueSize + 1;
     }
 
     /// @notice Get the number of pending registration requests in the waiting queue
     /// TODO modify this according to the current spec
-    /// TODO modify this according to the current spec
     function numPendingRegistrations() external view override returns (uint64) {
         return _numPendingRegistrations;
     }
 
     /// @notice Get the next available epoch for exit and queue size in that epoch
     /// TODO modify this according to the current spec
-    /// TODO modify this according to the current spec
     function nextExitEpoch() external view override returns (uint64, uint64) {
         uint64 epoch;
         uint64 queueSize;
@@ -212,15 +208,13 @@ contract StakeTable is AbstractStakeTable {
     // @param epoch next available exit epoch
     // @param queueSize current size of the exit queue (after insertion of new element in the queue)
     /// TODO modify this according to the current spec
-    /// TODO modify this according to the current spec
     function appendExitQueue(uint64 epoch, uint64 queueSize) private {
         firstAvailableExitEpoch = epoch;
         _numPendingExits = queueSize + 1;
     }
 
     /// @notice Get the number of pending exit requests in the waiting queue
     /// TODO modify this according to the current spec
-    /// TODO modify this according to the current spec
     function numPendingExits() external view override returns (uint64) {
         return _numPendingExits;
     }
@@ -240,7 +234,6 @@ contract StakeTable is AbstractStakeTable {
     /// @param node node which is assigned an exit escrow period.
     /// @return Number of epochs post exit after which funds can be withdrawn.
     /// TODO modify this according to the current spec
-    /// TODO modify this according to the current spec
     function exitEscrowPeriod(Node memory node) public pure returns (uint64) {
         if (node.balance > 100) {
             return 10;
@@ -402,7 +395,6 @@ contract StakeTable is AbstractStakeTable {
     function requestExit() external override {
         Node memory node = nodes[msg.sender];
 
-        // TODO test this behaviour when re-implementing the logic for handling epochs
         // if the node is not registered, revert
         if (node.account == address(0)) {
             revert NodeNotRegistered();
@@ -437,6 +429,8 @@ contract StakeTable is AbstractStakeTable {
     /// withdraw past their `exitEpoch`.
     ///
     /// @return The total amount withdrawn, equal to `Node.balance` associated with `blsVK`
+    /// TODO: add epoch logic so that we can ensure the node has first requested to exit and waiting
+    /// for the exit escrow period to be over
     function withdrawFunds() external override returns (uint256) {
         Node memory node = nodes[msg.sender];
 
@@ -456,10 +450,11 @@ contract StakeTable is AbstractStakeTable {
             revert InsufficientStakeBalance(0);
         }
 
-        // Verify that the exit escrow period is over.
-        if (currentEpoch() < node.exitEpoch + exitEscrowPeriod(node)) {
-            revert PrematureWithdrawal();
-        }
+        // // Verify that the exit escrow period is over.
+        // if (currentEpoch() < node.exitEpoch + exitEscrowPeriod(node)) {
+        //     revert PrematureWithdrawal();
+        // }
+        totalStake -= balance;
 
         // Delete the node from the stake table.
         delete nodes[msg.sender];

contracts/test/StakeTable.t.sol

@@ -673,4 +673,76 @@ contract StakeTable_register_Test is Test {
         );
         assertTrue(EdOnBN254.isEqual(node.schnorrVK, EdOnBN254.EdOnBN254Point(0, 0)));
     }
+
+    function test_WithdrawFunds_succeeds() public {
+        // Register the node and set exit epoch
+        uint64 depositAmount = 10 ether;
+        uint64 validUntilEpoch = 5;
+        string memory seed = "123";
+
+        (
+            BN254.G2Point memory blsVK,
+            EdOnBN254.EdOnBN254Point memory schnorrVK,
+            BN254.G1Point memory sig
+        ) = genClientWallet(exampleTokenCreator, seed);
+
+        // Prepare for the token transfer by granting allowance to the contract
+        vm.startPrank(exampleTokenCreator);
+        token.approve(address(stakeTable), depositAmount);
+
+        // Balances before registration
+        assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);
+
+        // register the node
+        vm.expectEmit(false, false, false, true, address(stakeTable));
+        emit AbstractStakeTable.Registered(exampleTokenCreator, 1, depositAmount);
+        stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);
+
+        // Withdraw the funds
+        uint256 balance = stakeTable.withdrawFunds();
+
+        // verify the withdraw
+        assertEq(balance, depositAmount);
+        assertEq(token.balanceOf(exampleTokenCreator), INITIAL_BALANCE);
+        assertEq(stakeTable.totalStake(), 0);
+        assertEq(stakeTable.lookupNode(exampleTokenCreator).balance, 0);
+        assertEq(stakeTable.lookupNode(exampleTokenCreator).account, address(0));
+
+        // test withdraw fails if the user tries to withdraw again
+        vm.expectRevert(S.NodeNotRegistered.selector);
+        stakeTable.withdrawFunds();
+
+        vm.stopPrank();
+    }
+
+    function test_WithdrawFunds_RevertWhen_NodeNotRegistered() public {
+        // Register the node and set exit epoch
+        uint64 depositAmount = 10 ether;
+        uint64 validUntilEpoch = 5;
+        string memory seed = "123";
+
+        // generate a new blsVK and schnorrVK and register this node
+        (
+            BN254.G2Point memory blsVK,
+            EdOnBN254.EdOnBN254Point memory schnorrVK,
+            BN254.G1Point memory sig
+        ) = genClientWallet(exampleTokenCreator, seed);
+
+        // Prepare for the token transfer by granting allowance to the contract
+        vm.startPrank(exampleTokenCreator);
+        token.approve(address(stakeTable), depositAmount);
+
+        // register the node
+        vm.expectEmit(false, false, false, true, address(stakeTable));
+        emit AbstractStakeTable.Registered(exampleTokenCreator, 1, depositAmount);
+        stakeTable.register(blsVK, schnorrVK, depositAmount, sig, validUntilEpoch);
+
+        vm.stopPrank();
+
+        vm.startPrank(makeAddr("randomUser"));
+        // withdraw the funds
+        vm.expectRevert(S.NodeNotRegistered.selector);
+        stakeTable.withdrawFunds();
+        vm.stopPrank();
+    }
 }