add withdrawFunds tests, removed epoch logic until from withdrawFunds…

… function until epochs are implemented, removed unecessary checks from the withdrawFunds function

contracts/src/StakeTable.sol

@@ -408,6 +408,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];
 
@@ -416,21 +418,17 @@ contract StakeTable is AbstractStakeTable {
             revert NodeNotRegistered();
         }
 
-        // The exit request must come from the node's withdrawal account.
-        if (node.account != msg.sender) {
-            revert Unauthenticated();
-        }
-
         // Verify that the balance is greater than zero
         uint256 balance = node.balance;
         if (balance == 0) {
             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

@@ -593,4 +593,72 @@ 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));
+
+        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();
+    }
 }