fix: unnaccounted for gas costs in withdraw tests established and set

Author: quirps

test/TokenVault.t.sol

@@ -5,14 +5,16 @@ import {TokenVault} from "../src/TokenVault.sol";
 contract TokenVaultTest is Test {
     TokenVault public vault;
     address public user1 = address(0xBEEF);
+    address public user2 = address(0xBEEF11);
     uint256 public constant INITIAL_BALANCE = 10 ether;
-
+    uint256 public constant WITHDRAW_AMOUNT = 0.1 ether;
     function setUp() public {
         // 1. Deploy the contract
         vault = new TokenVault();
 
         // 2. Fund the test account (user1) for transactions
         vm.deal(user1, INITIAL_BALANCE);
+        vm.deal(user2, INITIAL_BALANCE);
     }
 
     function testDeposit() public {
@@ -72,32 +74,48 @@ contract TokenVaultTest is Test {
     }
 
     // Invariant: The totalValue must always equal the sum of all user balances
-function invariant_TotalValueEqualsSumOfBalances() public view {
-    // This invariant function is run after every random call in the sequence.
-    uint256 sumOfBalances = 
-        vault.balances(user1) + 
-        vault.balances(address(0xDEAD)); // Imagine tracking a few key addresses
+    function invariant_TotalValueEqualsSumOfBalances() public view {
+        // This invariant function is run after every random call in the sequence.
+        uint256 sumOfBalances = vault.balances(user1) +
+            vault.balances(address(0xDEAD)); // Imagine tracking a few key addresses
 
-    assertEq(vault.totalValue(), sumOfBalances, "Invariant violated: totalValue != sum of balances");
-}
+        assertEq(
+            vault.totalValue(),
+            sumOfBalances,
+            "Invariant violated: totalValue != sum of balances"
+        );
+    }
 
-function testCompareWithdrawGas() public {
-    uint256 depositAmount = 1 ether;
-    uint256 withdrawAmount = 0.1 ether;
-
-    vm.deal(user1, depositAmount);
-    vm.prank(user1);
-    vault.deposit{value: depositAmount}(depositAmount);
-    
-    // 1. Call the optimized version
-    vm.prank(user1);
-    vault.withdrawOptimized(withdrawAmount);
-    
-    // 2. Call the unoptimized version (requires resetting the state, usually in a separate test)
-    // To properly compare, you should call them in separate test functions:
-    // function testWithdrawOptimized() public { /* ... call vault.withdrawOptimized ... */ }
-    // function testWithdrawUnoptimized() public { /* ... call vault.withdrawUnoptimized ... */ }
-    
-    // In this example, we'll just run the report on all functions.
-}
+    function testCompareWithdrawGas_Optimized() public {
+        // This test runs completely independently of the unoptimized one,
+        // ensuring a fresh state (cold storage) for the withdrawal.
+
+        // We're withdrawing only part of the deposit, so the vault's state
+        // (e.g., balance, shares) will be written to for the first time
+        // in this specific test function's execution.
+        uint256 preUserBalance = vault.balances(user1);
+        uint256 amount = 1 ether;
+        // 2. Execute with random 'amount'
+        vm.startPrank(user1);
+        vault.deposit{value: amount}(amount);
+        vault.withdrawOptimized(WITHDRAW_AMOUNT);
+        vm.stopPrank();
+
+        // Note: No need for assertions here unless you are specifically testing
+        // correctness. The primary goal is the gas report.
+    }
+
+    // 2. Test for the Unoptimized version
+    function testCompareWithdrawGas_Unoptimized() public {
+        // This test also runs independently, ensuring a fresh state (cold storage)
+        // for the unoptimized withdrawal.
+        uint256 preUserBalance = vault.balances(user2);
+        uint256 amount = 1 ether;
+        // 2. Execute with random 'amount'
+        vm.startPrank(user2);
+        vault.deposit{value: amount}(amount);
+        // If your Vault contract has an unoptimized function named withdrawUnoptimized:
+        vault.withdrawUnoptimized(WITHDRAW_AMOUNT);
+        vm.stopPrank();
+    }
 }