From 33781a53d04255f961707d3cdf35ad3fbf27d73c Mon Sep 17 00:00:00 2001
From: Ricardo Guilherme Schmidt <3esmit@gmail.com>
Date: Thu, 21 Aug 2025 13:19:16 -0300
Subject: [PATCH] chore(test): add test for withdraw after leave

---
 test/RewardsStreamerMP.t.sol | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol
index 7404dd5..100b6a6 100644
--- a/test/RewardsStreamerMP.t.sol
+++ b/test/RewardsStreamerMP.t.sol
@@ -2353,6 +2353,33 @@ contract LeaveTest is StakeManagerTest {
 
         assertEq(stakingToken.balanceOf(alice), aliceInitialBalance, "Alice has withdrawn her funds");
     }
+
+    
+    function test_WithdrawFromVaultCannotHappenBeforeLockupEnd() public {
+        uint256 aliceInitialBalance = stakingToken.balanceOf(alice);
+        uint256 stakeAmount = 10e18;
+        uint256 lockUpPeriod = streamer.MIN_LOCKUP_PERIOD();
+        _stake(alice, stakeAmount, lockUpPeriod);
+        assertEq(stakingToken.balanceOf(alice), aliceInitialBalance - stakeAmount, "Alice should have staked tokens");
+        _leave(alice);
+        vm.startPrank(alice);
+        vm.expectRevert(StakeVault.StakeVault__FundsLocked.selector);
+        StakeVault(vaults[alice]).withdrawFromVault(stakeAmount, alice);
+        assertEq(stakingToken.balanceOf(alice), aliceInitialBalance - stakeAmount, "Alice should still not have her funds back");
+    }
+
+    function test_WithdrawStakingTokenCannotHappenBeforeLockupEnd() public {
+        uint256 aliceInitialBalance = stakingToken.balanceOf(alice);
+        uint256 stakeAmount = 10e18;
+        uint256 lockUpPeriod = streamer.MIN_LOCKUP_PERIOD();
+        _stake(alice, stakeAmount, lockUpPeriod);
+        assertEq(stakingToken.balanceOf(alice), aliceInitialBalance - stakeAmount, "Alice should have staked tokens");
+        _leave(alice);
+        vm.prank(alice);
+        StakeVault(vaults[alice]).withdraw(stakingToken, stakeAmount);
+        assertEq(stakingToken.balanceOf(alice), aliceInitialBalance - stakeAmount, "Alice should still not have her funds back");
+    }
+
 }
 
 contract MaliciousUpgradeTest is StakeManagerTest {