From a2dfea656320760735be20779783a7ef443e45f7 Mon Sep 17 00:00:00 2001 From: luhrhenz Date: Mon, 2 Mar 2026 13:19:41 +0100 Subject: [PATCH] Add comprehensive Foundry tests for TimeLock and CounterV3 - Complete TimeLockV1 tests (30 tests, all passing) - Improved CounterV3 tests (30 tests, all passing) - Added Before->Action->After pattern for state changes - Added event testing for all events - All tests follow Foundry best practices --- assignments/foundry test/src/CounterV3.sol | 78 ++ assignments/foundry test/src/TimeLock.sol | 152 ++++ assignments/foundry test/test/CounterV3.t.sol | 385 +++++++++ .../foundry test/test/TimeLockV1.t.sol | 788 ++++++++++++++++++ 4 files changed, 1403 insertions(+) create mode 100644 assignments/foundry test/src/CounterV3.sol create mode 100644 assignments/foundry test/src/TimeLock.sol create mode 100644 assignments/foundry test/test/CounterV3.t.sol create mode 100644 assignments/foundry test/test/TimeLockV1.t.sol diff --git a/assignments/foundry test/src/CounterV3.sol b/assignments/foundry test/src/CounterV3.sol new file mode 100644 index 00000000..4a659183 --- /dev/null +++ b/assignments/foundry test/src/CounterV3.sol @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.28; + +/** + * @title CounterV3 + * @dev Version 3: Owner-only access + Privilege system + * Owner can grant/revoke privilege to non-owners to call state-changing functions + */ +contract CounterV3 { + uint public x; + address public owner; + mapping(address => bool) public privileged; + + event Increment(uint by); + event Decrement(uint by); + event PrivilegeGranted(address indexed account); + event PrivilegeRevoked(address indexed account); + + modifier onlyOwner() { + require(msg.sender == owner, "CounterV3: caller is not the owner"); + _; + } + + modifier onlyOwnerOrPrivileged() { + require( + msg.sender == owner || privileged[msg.sender], + "CounterV3: caller is not owner or privileged" + ); + _; + } + + constructor() { + owner = msg.sender; + } + + function inc() public onlyOwnerOrPrivileged { + x++; + emit Increment(1); + } + + function incBy(uint by) public onlyOwnerOrPrivileged { + require(by > 0, "incBy: increment should be positive"); + x += by; + emit Increment(by); + } + + /** + * @dev Decreases the counter by the specified amount + * @param by The amount to decrease (must be positive and not exceed current counter value) + */ + function decrease(uint by) public onlyOwnerOrPrivileged { + require(by > 0, "decrease: decrement should be positive"); + require(x >= by, "decrease: counter cannot go below zero"); + x -= by; + emit Decrement(by); + } + + /** + * @dev Grants privilege to a non-owner to call state-changing functions + * @param account The address to grant privilege to + */ + function grantPrivilege(address account) public onlyOwner { + require(account != address(0), "CounterV3: invalid address"); + require(account != owner, "CounterV3: owner is already privileged"); + privileged[account] = true; + emit PrivilegeGranted(account); + } + + /** + * @dev Revokes privilege from a non-owner + * @param account The address to revoke privilege from + */ + function revokePrivilege(address account) public onlyOwner { + require(privileged[account], "CounterV3: account is not privileged"); + privileged[account] = false; + emit PrivilegeRevoked(account); + } +} diff --git a/assignments/foundry test/src/TimeLock.sol b/assignments/foundry test/src/TimeLock.sol new file mode 100644 index 00000000..59614865 --- /dev/null +++ b/assignments/foundry test/src/TimeLock.sol @@ -0,0 +1,152 @@ +//SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +contract TimeLockV1 { + + struct Vault { + uint balance; + uint unlockTime; + bool active; + } + + mapping(address => Vault[]) private vaults; + + event Deposited(address indexed user, uint vaultId, uint amount, uint unlockTime); + event Withdrawn(address indexed user, uint vaultId, uint amount); + + function deposit(uint _unlockTime) external payable returns (uint) { + require(msg.value > 0, "Deposit must be greater than zero"); + require(_unlockTime > block.timestamp, "Unlock time must be in the future"); + + // Create new vault + vaults[msg.sender].push(Vault({ + balance: msg.value, + unlockTime: _unlockTime, + active: true + })); + + uint vaultId = vaults[msg.sender].length - 1; + emit Deposited(msg.sender, vaultId, msg.value, _unlockTime); + + return vaultId; + } + function withdraw(uint _vaultId) external { + require(_vaultId < vaults[msg.sender].length, "Invalid vault ID"); + + Vault storage userVault = vaults[msg.sender][_vaultId]; + require(userVault.active, "Vault is not active"); + require(userVault.balance > 0, "Vault has zero balance"); + require(block.timestamp >= userVault.unlockTime, "Funds are still locked"); + + uint amount = userVault.balance; + + // Mark vault as inactive and clear balance + userVault.balance = 0; + userVault.active = false; + + (bool success, ) = payable(msg.sender).call{value: amount}(""); + require(success, "Transfer failed"); + + emit Withdrawn(msg.sender, _vaultId, amount); + } + function withdrawAll() external returns (uint) { + uint totalWithdrawn = 0; + Vault[] storage userVaults = vaults[msg.sender]; + + for (uint i = 0; i < userVaults.length; i++) { + if (userVaults[i].active && + userVaults[i].balance > 0 && + block.timestamp >= userVaults[i].unlockTime) { + + uint amount = userVaults[i].balance; + userVaults[i].balance = 0; + userVaults[i].active = false; + + totalWithdrawn += amount; + emit Withdrawn(msg.sender, i, amount); + } + } + + require(totalWithdrawn > 0, "No unlocked funds available"); + + (bool success, ) = payable(msg.sender).call{value: totalWithdrawn}(""); + require(success, "Transfer failed"); + + return totalWithdrawn; + } + function getVaultCount(address _user) external view returns (uint) { + return vaults[_user].length; + } + function getVault(address _user, uint _vaultId) external view returns ( + uint balance, + uint unlockTime, + bool active, + bool isUnlocked + ) { + require(_vaultId < vaults[_user].length, "Invalid vault ID"); + + Vault storage vault = vaults[_user][_vaultId]; + return ( + vault.balance, + vault.unlockTime, + vault.active, + block.timestamp >= vault.unlockTime + ); + } + function getAllVaults(address _user) external view returns (Vault[] memory) { + return vaults[_user]; + } + function getActiveVaults(address _user) external view returns ( + uint[] memory activeVaults, + uint[] memory balances, + uint[] memory unlockTimes + ) { + Vault[] storage userVaults = vaults[_user]; + + // Count active vaults + uint activeCount = 0; + for (uint i = 0; i < userVaults.length; i++) { + if (userVaults[i].active && userVaults[i].balance > 0) { + activeCount++; + } + } + + // Create arrays + activeVaults = new uint[](activeCount); + balances = new uint[](activeCount); + unlockTimes = new uint[](activeCount); + + // Populate arrays + uint index = 0; + for (uint i = 0; i < userVaults.length; i++) { + if (userVaults[i].active && userVaults[i].balance > 0) { + activeVaults[index] = i; + balances[index] = userVaults[i].balance; + unlockTimes[index] = userVaults[i].unlockTime; + index++; + } + } + + return (activeVaults, balances, unlockTimes); + } + function getTotalBalance(address _user) external view returns (uint total) { + Vault[] storage userVaults = vaults[_user]; + for (uint i = 0; i < userVaults.length; i++) { + if (userVaults[i].active) { + total += userVaults[i].balance; + } + } + return total; + } + function getUnlockedBalance(address _user) external view returns (uint unlocked) { + Vault[] storage userVaults = vaults[_user]; + for (uint i = 0; i < userVaults.length; i++) { + if (userVaults[i].active && + userVaults[i].balance > 0 && + block.timestamp >= userVaults[i].unlockTime) { + unlocked += userVaults[i].balance; + } + } + return unlocked; + } +} diff --git a/assignments/foundry test/test/CounterV3.t.sol b/assignments/foundry test/test/CounterV3.t.sol new file mode 100644 index 00000000..9fd55240 --- /dev/null +++ b/assignments/foundry test/test/CounterV3.t.sol @@ -0,0 +1,385 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.28; + +import {CounterV3} from "../src/CounterV3.sol"; +import {Test} from "forge-std/Test.sol"; + +contract CounterV3Test is Test { + CounterV3 counter; + address owner; + address privilegedUser; + address nonPrivilegedUser; + + // Events to test + event Increment(uint by); + event Decrement(uint by); + event PrivilegeGranted(address indexed account); + event PrivilegeRevoked(address indexed account); + + function setUp() public { + owner = address(this); + privilegedUser = address(0x123); + nonPrivilegedUser = address(0x456); + + counter = new CounterV3(); + } + + // Test 1: Deployment and initial state + function test_InitialState() public view { + assertEq(counter.x(), 0, "Initial counter value should be 0"); + assertEq(counter.owner(), owner, "Deployer should be owner"); + assertFalse(counter.privileged(privilegedUser), "New address should not be privileged"); + assertFalse(counter.privileged(nonPrivilegedUser), "Non-privileged user should not be privileged"); + } + + // Test 2: Owner can increment + function test_OwnerCanIncrement() public { + // Check state BEFORE + uint256 counterBefore = counter.x(); + assertEq(counterBefore, 0, "Counter should start at 0"); + + // Perform action + counter.inc(); + + // Check state AFTER + uint256 counterAfter = counter.x(); + assertEq(counterAfter, 1, "Counter should increment to 1"); + assertEq(counterAfter - counterBefore, 1, "Counter should increase by 1"); + } + + // Test 3: Non-owner cannot increment without privilege + function test_NonOwnerCannotIncrementWithoutPrivilege() public { + vm.prank(nonPrivilegedUser); + vm.expectRevert("CounterV3: caller is not owner or privileged"); + counter.inc(); + } + + // Test 4: Owner can grant privilege + function test_OwnerCanGrantPrivilege() public { + // Check state BEFORE + bool privilegedBefore = counter.privileged(privilegedUser); + assertFalse(privilegedBefore, "User should not be privileged initially"); + + // Perform action + counter.grantPrivilege(privilegedUser); + + // Check state AFTER + bool privilegedAfter = counter.privileged(privilegedUser); + assertTrue(privilegedAfter, "User should be privileged after grant"); + } + + // Test 5: Non-owner cannot grant privilege + function test_NonOwnerCannotGrantPrivilege() public { + vm.prank(nonPrivilegedUser); + vm.expectRevert("CounterV3: caller is not the owner"); + counter.grantPrivilege(privilegedUser); + } + + // Test 6: Cannot grant privilege to zero address + function test_CannotGrantPrivilegeToZeroAddress() public { + vm.expectRevert("CounterV3: invalid address"); + counter.grantPrivilege(address(0)); + } + + // Test 7: Cannot grant privilege to owner + function test_CannotGrantPrivilegeToOwner() public { + vm.expectRevert("CounterV3: owner is already privileged"); + counter.grantPrivilege(owner); + } + + // Test 8: Privileged user can increment + function test_PrivilegedUserCanIncrement() public { + // Setup: Grant privilege + counter.grantPrivilege(privilegedUser); + + // Check state BEFORE + uint256 counterBefore = counter.x(); + assertEq(counterBefore, 0, "Counter should start at 0"); + + // Perform action + vm.prank(privilegedUser); + counter.inc(); + + // Check state AFTER + uint256 counterAfter = counter.x(); + assertEq(counterAfter, 1, "Privileged user should be able to increment"); + assertEq(counterAfter - counterBefore, 1, "Counter should increase by 1"); + } + + // Test 9: Owner can revoke privilege + function test_OwnerCanRevokePrivilege() public { + // Setup: Grant privilege first + counter.grantPrivilege(privilegedUser); + + // Check state BEFORE revoke + bool privilegedBefore = counter.privileged(privilegedUser); + assertTrue(privilegedBefore, "User should be privileged before revoke"); + + // Perform action + counter.revokePrivilege(privilegedUser); + + // Check state AFTER revoke + bool privilegedAfter = counter.privileged(privilegedUser); + assertFalse(privilegedAfter, "User should not be privileged after revoke"); + } + + // Test 10: Non-owner cannot revoke privilege + function test_NonOwnerCannotRevokePrivilege() public { + counter.grantPrivilege(privilegedUser); + + vm.prank(nonPrivilegedUser); + vm.expectRevert("CounterV3: caller is not the owner"); + counter.revokePrivilege(privilegedUser); + } + + // Test 11: Cannot revoke privilege from non-privileged account + function test_CannotRevokePrivilegeFromNonPrivileged() public { + vm.expectRevert("CounterV3: account is not privileged"); + counter.revokePrivilege(privilegedUser); + } + + // Test 12: Revoked user cannot increment + function test_RevokedUserCannotIncrement() public { + counter.grantPrivilege(privilegedUser); + counter.revokePrivilege(privilegedUser); + + vm.prank(privilegedUser); + vm.expectRevert("CounterV3: caller is not owner or privileged"); + counter.inc(); + } + + // Test 13: incBy function works for owner + function test_OwnerCanIncrementBy() public { + // Check state BEFORE + uint256 counterBefore = counter.x(); + assertEq(counterBefore, 0, "Counter should start at 0"); + + // Perform action + counter.incBy(5); + + // Check state AFTER + uint256 counterAfter = counter.x(); + assertEq(counterAfter, 5, "Counter should increment by 5"); + assertEq(counterAfter - counterBefore, 5, "Counter should increase by exactly 5"); + } + + // Test 14: incBy function works for privileged user + function test_PrivilegedUserCanIncrementBy() public { + // Setup: Grant privilege + counter.grantPrivilege(privilegedUser); + + // Check state BEFORE + uint256 counterBefore = counter.x(); + assertEq(counterBefore, 0, "Counter should start at 0"); + + // Perform action + vm.prank(privilegedUser); + counter.incBy(10); + + // Check state AFTER + uint256 counterAfter = counter.x(); + assertEq(counterAfter, 10, "Privileged user should increment by 10"); + assertEq(counterAfter - counterBefore, 10, "Counter should increase by exactly 10"); + } + + // Test 15: incBy requires positive value + function test_incByRequiresPositiveValue() public { + vm.expectRevert("incBy: increment should be positive"); + counter.incBy(0); + } + + // Test 16: decrease function works for owner + function test_OwnerCanDecrease() public { + // Setup: Increment first + counter.incBy(10); + + // Check state BEFORE decrease + uint256 counterBefore = counter.x(); + assertEq(counterBefore, 10, "Counter should be at 10"); + + // Perform action + counter.decrease(3); + + // Check state AFTER decrease + uint256 counterAfter = counter.x(); + assertEq(counterAfter, 7, "Counter should decrease to 7"); + assertEq(counterBefore - counterAfter, 3, "Counter should decrease by exactly 3"); + } + + // Test 17: decrease function works for privileged user + function test_PrivilegedUserCanDecrease() public { + // Setup: Increment and grant privilege + counter.incBy(15); + counter.grantPrivilege(privilegedUser); + + // Check state BEFORE decrease + uint256 counterBefore = counter.x(); + assertEq(counterBefore, 15, "Counter should be at 15"); + + // Perform action + vm.prank(privilegedUser); + counter.decrease(5); + + // Check state AFTER decrease + uint256 counterAfter = counter.x(); + assertEq(counterAfter, 10, "Privileged user should decrease by 5"); + assertEq(counterBefore - counterAfter, 5, "Counter should decrease by exactly 5"); + } + + // Test 18: decrease requires positive value + function test_decreaseRequiresPositiveValue() public { + vm.expectRevert("decrease: decrement should be positive"); + counter.decrease(0); + } + + // Test 19: decrease cannot go below zero + function test_decreaseCannotGoBelowZero() public { + counter.incBy(5); + + vm.expectRevert("decrease: counter cannot go below zero"); + counter.decrease(10); + } + + // Test 20: Multiple privileged users + function test_MultiplePrivilegedUsers() public { + address privilegedUser2 = address(0x789); + + // Check state BEFORE + uint256 counterBefore = counter.x(); + assertEq(counterBefore, 0, "Counter should start at 0"); + + // Setup: Grant privileges + counter.grantPrivilege(privilegedUser); + counter.grantPrivilege(privilegedUser2); + + // Perform actions + vm.prank(privilegedUser); + counter.inc(); + + vm.prank(privilegedUser2); + counter.incBy(3); + + // Check state AFTER + uint256 counterAfter = counter.x(); + assertEq(counterAfter, 4, "Multiple privileged users should work"); + assertEq(counterAfter - counterBefore, 4, "Counter should increase by 4 total"); + } + + // 🎯 ADDITIONAL TESTS FOR EVENTS AND COMPLETE COVERAGE + + // Test 21: inc() emits Increment event + function test_IncEmitsEvent() public { + // Expect event + vm.expectEmit(true, true, true, true); + emit Increment(1); + + // Perform action + counter.inc(); + } + + // Test 22: incBy() emits Increment event + function test_IncByEmitsEvent() public { + // Expect event + vm.expectEmit(true, true, true, true); + emit Increment(5); + + // Perform action + counter.incBy(5); + } + + // Test 23: decrease() emits Decrement event + function test_DecreaseEmitsEvent() public { + // Setup + counter.incBy(10); + + // Expect event + vm.expectEmit(true, true, true, true); + emit Decrement(3); + + // Perform action + counter.decrease(3); + } + + // Test 24: grantPrivilege() emits PrivilegeGranted event + function test_GrantPrivilegeEmitsEvent() public { + // Expect event + vm.expectEmit(true, false, false, true); + emit PrivilegeGranted(privilegedUser); + + // Perform action + counter.grantPrivilege(privilegedUser); + } + + // Test 25: revokePrivilege() emits PrivilegeRevoked event + function test_RevokePrivilegeEmitsEvent() public { + // Setup + counter.grantPrivilege(privilegedUser); + + // Expect event + vm.expectEmit(true, false, false, true); + emit PrivilegeRevoked(privilegedUser); + + // Perform action + counter.revokePrivilege(privilegedUser); + } + + // Test 26: Non-privileged user cannot use incBy + function test_NonPrivilegedUserCannotIncBy() public { + vm.prank(nonPrivilegedUser); + vm.expectRevert("CounterV3: caller is not owner or privileged"); + counter.incBy(5); + } + + // Test 27: Non-privileged user cannot use decrease + function test_NonPrivilegedUserCannotDecrease() public { + counter.incBy(10); + + vm.prank(nonPrivilegedUser); + vm.expectRevert("CounterV3: caller is not owner or privileged"); + counter.decrease(5); + } + + // Test 28: Counter state persists across multiple operations + function test_CounterStatePersistence() public { + // Check initial state + assertEq(counter.x(), 0, "Should start at 0"); + + // Multiple operations + counter.inc(); + assertEq(counter.x(), 1, "Should be 1 after inc"); + + counter.incBy(5); + assertEq(counter.x(), 6, "Should be 6 after incBy(5)"); + + counter.decrease(2); + assertEq(counter.x(), 4, "Should be 4 after decrease(2)"); + + counter.inc(); + assertEq(counter.x(), 5, "Should be 5 after final inc"); + } + + // Test 29: Privilege state persists correctly + function test_PrivilegeStatePersistence() public { + // Initially not privileged + assertFalse(counter.privileged(privilegedUser), "Should not be privileged initially"); + + // Grant privilege + counter.grantPrivilege(privilegedUser); + assertTrue(counter.privileged(privilegedUser), "Should be privileged after grant"); + + // Privilege persists + vm.prank(privilegedUser); + counter.inc(); + assertTrue(counter.privileged(privilegedUser), "Should still be privileged after using it"); + + // Revoke privilege + counter.revokePrivilege(privilegedUser); + assertFalse(counter.privileged(privilegedUser), "Should not be privileged after revoke"); + } + + // Test 30: Owner remains owner throughout + function test_OwnershipPersistence() public view { + assertEq(counter.owner(), owner, "Owner should be set correctly"); + // Owner never changes in this contract (no transfer function) + } +} \ No newline at end of file diff --git a/assignments/foundry test/test/TimeLockV1.t.sol b/assignments/foundry test/test/TimeLockV1.t.sol new file mode 100644 index 00000000..e08dd05c --- /dev/null +++ b/assignments/foundry test/test/TimeLockV1.t.sol @@ -0,0 +1,788 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +import "forge-std/Test.sol"; +import "../src/TimeLock.sol"; + +/** + * 🎮 FOUNDRY TESTING QUEST 🎮 + * + * Welcome, brave tester! Your mission is to write Solidity tests for the TimeLockV1 contract. + * + * 📚 FOUNDRY BASICS: + * - Test functions must start with "test" + * - Use setUp() to initialize before each test + * - Use vm.prank(address) to simulate calls from different addresses + * - Use vm.warp(timestamp) to change block.timestamp + * - Use vm.deal(address, amount) to give ETH to an address + * - Use vm.expectRevert() to test for reverts + * + * 🔧 USEFUL ASSERTIONS: + * - assertEq(a, b) - assert a equals b + * - assertTrue(condition) - assert condition is true + * - assertFalse(condition) - assert condition is false + * - assertGt(a, b) - assert a > b + * - assertLt(a, b) - assert a < b + * + * 🎯 YOUR QUESTS: + * + * QUEST 1: Setup & Deployment ⚔️ + * - Deploy the TimeLockV1 contract in setUp() + * - Create test addresses (user1, user2) + * - Give them some ETH to work with + * + * QUEST 2: Test Deposit Function 💰 + * - Test successful deposit + * - Test deposit with 0 value (should revert) + * - Test deposit with past unlock time (should revert) + * - Verify vault is created correctly + * - Check that Deposited event is emitted + * + * QUEST 3: Test Withdraw Function 🏦 + * - Test successful withdrawal after unlock time + * - Test withdrawal before unlock time (should revert) + * - Test withdrawal of non-existent vault (should revert) + * - Test that another user can't withdraw your vault + * - Verify balances change correctly + * + * QUEST 4: Test WithdrawAll Function 🌟 + * - Create multiple vaults with different unlock times + * - Test withdrawing all unlocked vaults at once + * - Verify only unlocked vaults are withdrawn + * - Test with no unlocked vaults (should revert) + * + * QUEST 5: Test View Functions 🔍 + * - Test getTotalBalance() + * - Test getUnlockedBalance() + * - Test getActiveVaults() + * - Test getVault() + * - Test getAllVaults() + * + * 💡 TIPS: + * - Start simple, then add complexity + * - Test one thing at a time + * - Use descriptive test names + * - Run tests with: forge test + * - Run specific test: forge test --match-test testDeposit + * - See gas usage: forge test --gas-report + * - Verbose output: forge test -vvvv + * + * Good luck, adventurer! 🚀 + */ + +contract TimeLockV1Test is Test { + TimeLockV1 public timelock; + + address public user1; + address public user2; + + // Events to test + event Deposited(address indexed user, uint vaultId, uint amount, uint unlockTime); + event Withdrawn(address indexed user, uint vaultId, uint amount); + + function setUp() public { + // 🎯 QUEST 1: Deploy and setup + timelock = new TimeLockV1(); + + // Create user addresses + user1 = address(1); + user2 = address(2); + + // Give users some ETH + vm.deal(user1, 1000 ether); + vm.deal(user2, 1000 ether); + } + + // 🎯 QUEST 2: DEPOSIT TESTS + + function testDeposit() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 1 ether; + + // Check initial state BEFORE deposit + uint256 userBalanceBefore = user1.balance; + uint256 contractBalanceBefore = address(timelock).balance; + uint256 vaultCountBefore = timelock.getVaultCount(user1); + + assertEq(userBalanceBefore, 1000 ether, "User should start with 1000 ether"); + assertEq(contractBalanceBefore, 0, "Contract should start with 0 balance"); + assertEq(vaultCountBefore, 0, "User should have 0 vaults initially"); + + // Perform action: deposit + vm.prank(user1); + uint256 vaultId = timelock.deposit{value: depositAmount}(unlockTime); + + // Check state AFTER deposit + uint256 userBalanceAfter = user1.balance; + uint256 contractBalanceAfter = address(timelock).balance; + uint256 vaultCountAfter = timelock.getVaultCount(user1); + + // Check vault was created correctly + (uint256 balance, uint256 unlock, bool active, bool isUnlocked) = + timelock.getVault(user1, 0); + + // Assert all state changes + assertEq(vaultId, 0, "First vault should have ID 0"); + assertEq(vaultCountAfter, 1, "User should have 1 vault after deposit"); + assertEq(userBalanceBefore - userBalanceAfter, depositAmount, "User balance should decrease by deposit amount"); + assertEq(contractBalanceAfter - contractBalanceBefore, depositAmount, "Contract balance should increase by deposit amount"); + assertEq(balance, depositAmount, "Vault balance should match deposit amount"); + assertEq(unlock, unlockTime, "Unlock time should match"); + assertTrue(active, "Vault should be active"); + assertFalse(isUnlocked, "Vault should not be unlocked yet"); + } + + function testDepositRevertsWithZeroValue() public { + uint256 unlockTime = block.timestamp + 1 hours; + + vm.prank(user1); + vm.expectRevert("Deposit must be greater than zero"); + timelock.deposit{value: 0}(unlockTime); + } + + function testDepositRevertsWithPastUnlockTime() public { + // Use a timestamp that's definitely in the past (timestamp 1) + uint256 pastTime = 1; + + vm.prank(user1); + vm.expectRevert("Unlock time must be in the future"); + timelock.deposit{value: 1 ether}(pastTime); + } + + // 🎯 QUEST 3: WITHDRAW TESTS + + function testWithdraw() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 1 ether; + + // Setup: Deposit first + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Move time forward past unlock time + vm.warp(unlockTime + 1); + + // Check state BEFORE withdrawal + uint256 userBalanceBefore = user1.balance; + uint256 contractBalanceBefore = address(timelock).balance; + (uint256 vaultBalanceBefore,, bool activeBefore, bool isUnlockedBefore) = + timelock.getVault(user1, 0); + + assertEq(userBalanceBefore, 999 ether, "User should have 999 ether after deposit"); + assertEq(contractBalanceBefore, 1 ether, "Contract should have 1 ether"); + assertEq(vaultBalanceBefore, depositAmount, "Vault should have deposit amount"); + assertTrue(activeBefore, "Vault should be active before withdrawal"); + assertTrue(isUnlockedBefore, "Vault should be unlocked"); + + // Perform action: Withdraw + vm.prank(user1); + timelock.withdraw(0); + + // Check state AFTER withdrawal + uint256 userBalanceAfter = user1.balance; + uint256 contractBalanceAfter = address(timelock).balance; + (uint256 vaultBalanceAfter,, bool activeAfter,) = timelock.getVault(user1, 0); + + // Assert all state changes + assertEq(vaultBalanceAfter, 0, "Vault balance should be 0 after withdrawal"); + assertFalse(activeAfter, "Vault should be inactive after withdrawal"); + assertEq(userBalanceAfter - userBalanceBefore, depositAmount, "User should receive deposit amount"); + assertEq(contractBalanceBefore - contractBalanceAfter, depositAmount, "Contract balance should decrease"); + assertEq(userBalanceAfter, 1000 ether, "User should have original 1000 ether back"); + assertEq(contractBalanceAfter, 0, "Contract should have 0 balance"); + } + + function testWithdrawRevertsBeforeUnlockTime() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // Deposit + vm.prank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + + // Try to withdraw before unlock time (don't warp time) + vm.prank(user1); + vm.expectRevert("Funds are still locked"); + timelock.withdraw(0); + } + + function testWithdrawRevertsForInvalidVaultId() public { + // Try to withdraw a vault that doesn't exist + vm.prank(user1); + vm.expectRevert("Invalid vault ID"); + timelock.withdraw(0); + } + + function testCannotWithdrawOtherUsersVault() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // user1 deposits + vm.prank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + + // Move time forward + vm.warp(unlockTime + 1); + + // user2 tries to withdraw user1's vault + // user2 has no vaults, so vault ID 0 is invalid for them + vm.prank(user2); + vm.expectRevert("Invalid vault ID"); + timelock.withdraw(0); + } + + // 🎯 QUEST 4: WITHDRAW ALL TESTS + + function testWithdrawAll() public { + // Create 3 vaults with different unlock times + uint256 unlockTime1 = block.timestamp + 1 hours; + uint256 unlockTime2 = block.timestamp + 2 hours; + uint256 unlockTime3 = block.timestamp + 10 hours; + + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime1); + timelock.deposit{value: 2 ether}(unlockTime2); + timelock.deposit{value: 3 ether}(unlockTime3); + vm.stopPrank(); + + // Move time to unlock first 2 vaults only + vm.warp(unlockTime2 + 1); + + // Check state BEFORE withdrawAll + uint256 userBalanceBefore = user1.balance; + uint256 contractBalanceBefore = address(timelock).balance; + uint256 totalBalanceBefore = timelock.getTotalBalance(user1); + uint256 unlockedBalanceBefore = timelock.getUnlockedBalance(user1); + (,, bool active1Before,) = timelock.getVault(user1, 0); + (,, bool active2Before,) = timelock.getVault(user1, 1); + (,, bool active3Before,) = timelock.getVault(user1, 2); + + assertEq(userBalanceBefore, 994 ether, "User should have 994 ether after 3 deposits"); + assertEq(contractBalanceBefore, 6 ether, "Contract should have 6 ether"); + assertEq(totalBalanceBefore, 6 ether, "Total balance should be 6 ether"); + assertEq(unlockedBalanceBefore, 3 ether, "Unlocked balance should be 3 ether (1+2)"); + assertTrue(active1Before, "Vault 1 should be active"); + assertTrue(active2Before, "Vault 2 should be active"); + assertTrue(active3Before, "Vault 3 should be active"); + + // Perform action: Withdraw all unlocked vaults + vm.prank(user1); + uint256 totalWithdrawn = timelock.withdrawAll(); + + // Check state AFTER withdrawAll + uint256 userBalanceAfter = user1.balance; + uint256 contractBalanceAfter = address(timelock).balance; + uint256 totalBalanceAfter = timelock.getTotalBalance(user1); + (,, bool active1After,) = timelock.getVault(user1, 0); + (,, bool active2After,) = timelock.getVault(user1, 1); + (,, bool active3After,) = timelock.getVault(user1, 2); + + // Assert all state changes + assertEq(totalWithdrawn, 3 ether, "Should withdraw 3 ether total"); + assertEq(userBalanceAfter - userBalanceBefore, 3 ether, "User should receive 3 ether"); + assertEq(contractBalanceBefore - contractBalanceAfter, 3 ether, "Contract should lose 3 ether"); + assertEq(userBalanceAfter, 997 ether, "User should have 997 ether"); + assertEq(contractBalanceAfter, 3 ether, "Contract should have 3 ether left"); + assertEq(totalBalanceAfter, 3 ether, "Total balance should be 3 ether (only vault 3)"); + assertFalse(active1After, "Vault 1 should be inactive"); + assertFalse(active2After, "Vault 2 should be inactive"); + assertTrue(active3After, "Vault 3 should still be active (locked)"); + } + + function testWithdrawAllRevertsWithNoUnlockedFunds() public { + // Create vaults that are still locked + uint256 unlockTime = block.timestamp + 10 hours; + + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + timelock.deposit{value: 2 ether}(unlockTime); + vm.stopPrank(); + + // Try to withdrawAll without moving time forward + vm.prank(user1); + vm.expectRevert("No unlocked funds available"); + timelock.withdrawAll(); + } + + // 🎯 QUEST 5: VIEW FUNCTION TESTS + + function testGetTotalBalance() public { + // Create multiple vaults + uint256 unlockTime = block.timestamp + 1 hours; + + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + timelock.deposit{value: 2 ether}(unlockTime); + timelock.deposit{value: 3 ether}(unlockTime); + vm.stopPrank(); + + // Check total balance + uint256 totalBalance = timelock.getTotalBalance(user1); + assertEq(totalBalance, 6 ether, "Total balance should be 6 ether"); + } + + function testGetUnlockedBalance() public { + // Create vaults with different unlock times + uint256 unlockTime1 = block.timestamp + 1 hours; + uint256 unlockTime2 = block.timestamp + 2 hours; + uint256 unlockTime3 = block.timestamp + 10 hours; + + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime1); + timelock.deposit{value: 2 ether}(unlockTime2); + timelock.deposit{value: 3 ether}(unlockTime3); + vm.stopPrank(); + + // Before any unlock + uint256 unlockedBefore = timelock.getUnlockedBalance(user1); + assertEq(unlockedBefore, 0, "No funds should be unlocked yet"); + + // Move time to unlock first 2 vaults + vm.warp(unlockTime2 + 1); + + // Check unlocked balance + uint256 unlockedAfter = timelock.getUnlockedBalance(user1); + assertEq(unlockedAfter, 3 ether, "Should have 3 ether unlocked (1 + 2)"); + } + + function testGetVault() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 5 ether; + + // Create a vault + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Get vault details + (uint256 balance, uint256 unlock, bool active, bool isUnlocked) = + timelock.getVault(user1, 0); + + assertEq(balance, depositAmount, "Balance should match deposit"); + assertEq(unlock, unlockTime, "Unlock time should match"); + assertTrue(active, "Vault should be active"); + assertFalse(isUnlocked, "Vault should not be unlocked yet"); + + // Move time forward + vm.warp(unlockTime + 1); + + // Check again + (,,,bool isUnlockedNow) = timelock.getVault(user1, 0); + assertTrue(isUnlockedNow, "Vault should be unlocked now"); + } + + // 🏆 BONUS QUESTS (Optional Challenges!) + + function testMultipleUsersCanHaveVaults() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // User1 deposits + vm.prank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + + // User2 deposits + vm.prank(user2); + timelock.deposit{value: 2 ether}(unlockTime); + + // Check user1's vault + (uint256 balance1,,,) = timelock.getVault(user1, 0); + assertEq(balance1, 1 ether, "User1 vault should have 1 ether"); + + // Check user2's vault + (uint256 balance2,,,) = timelock.getVault(user2, 0); + assertEq(balance2, 2 ether, "User2 vault should have 2 ether"); + + // Check vault counts + uint256 count1 = timelock.getVaultCount(user1); + uint256 count2 = timelock.getVaultCount(user2); + assertEq(count1, 1, "User1 should have 1 vault"); + assertEq(count2, 1, "User2 should have 1 vault"); + } + + function testDepositEmitsEvent() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 1 ether; + + // Expect event with these parameters + vm.expectEmit(true, false, false, true); + emit Deposited(user1, 0, depositAmount, unlockTime); + + // Make the call that should emit the event + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + } + + function testWithdrawEmitsEvent() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 1 ether; + + // Deposit first + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Move time forward + vm.warp(unlockTime + 1); + + // Expect withdraw event + vm.expectEmit(true, false, false, true); + emit Withdrawn(user1, 0, depositAmount); + + // Withdraw + vm.prank(user1); + timelock.withdraw(0); + } + + function testBalanceChangesOnDeposit() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 1 ether; + + // Check state BEFORE deposit + uint256 userBalanceBefore = user1.balance; + uint256 contractBalanceBefore = address(timelock).balance; + + assertEq(userBalanceBefore, 1000 ether, "User should start with 1000 ether"); + assertEq(contractBalanceBefore, 0, "Contract should start with 0 balance"); + + // Perform action: deposit + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Check state AFTER deposit + uint256 userBalanceAfter = user1.balance; + uint256 contractBalanceAfter = address(timelock).balance; + + // Assert state changes + assertEq(userBalanceBefore - userBalanceAfter, depositAmount, "User balance should decrease"); + assertEq(contractBalanceAfter - contractBalanceBefore, depositAmount, "Contract balance should increase"); + assertEq(userBalanceAfter, 999 ether, "User should have 999 ether"); + assertEq(contractBalanceAfter, 1 ether, "Contract should have 1 ether"); + } + + function testBalanceChangesOnWithdraw() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 1 ether; + + // Setup: Deposit first + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Move time forward + vm.warp(unlockTime + 1); + + // Check state BEFORE withdraw + uint256 userBalanceBefore = user1.balance; + uint256 contractBalanceBefore = address(timelock).balance; + + assertEq(userBalanceBefore, 999 ether, "User should have 999 ether after deposit"); + assertEq(contractBalanceBefore, 1 ether, "Contract should have 1 ether"); + + // Perform action: Withdraw + vm.prank(user1); + timelock.withdraw(0); + + // Check state AFTER withdraw + uint256 userBalanceAfter = user1.balance; + uint256 contractBalanceAfter = address(timelock).balance; + + // Assert state changes + assertEq(userBalanceAfter - userBalanceBefore, depositAmount, "User balance should increase"); + assertEq(contractBalanceBefore - contractBalanceAfter, depositAmount, "Contract balance should decrease"); + assertEq(userBalanceAfter, 1000 ether, "User should have 1000 ether back"); + assertEq(contractBalanceAfter, 0, "Contract should have 0 ether"); + } + + function testGetActiveVaults() public { + uint256 unlockTime1 = block.timestamp + 1 hours; + uint256 unlockTime2 = block.timestamp + 2 hours; + + // Create 2 vaults + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime1); + timelock.deposit{value: 2 ether}(unlockTime2); + vm.stopPrank(); + + // Get active vaults + (uint[] memory activeVaults, uint[] memory balances, uint[] memory unlockTimes) = + timelock.getActiveVaults(user1); + + assertEq(activeVaults.length, 2, "Should have 2 active vaults"); + assertEq(balances[0], 1 ether, "First vault should have 1 ether"); + assertEq(balances[1], 2 ether, "Second vault should have 2 ether"); + assertEq(unlockTimes[0], unlockTime1, "First unlock time should match"); + assertEq(unlockTimes[1], unlockTime2, "Second unlock time should match"); + } + + function testGetAllVaults() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // Create 2 vaults + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + timelock.deposit{value: 2 ether}(unlockTime); + vm.stopPrank(); + + // Get all vaults + TimeLockV1.Vault[] memory vaults = timelock.getAllVaults(user1); + + assertEq(vaults.length, 2, "Should have 2 vaults"); + assertEq(vaults[0].balance, 1 ether, "First vault balance"); + assertEq(vaults[1].balance, 2 ether, "Second vault balance"); + } + + function testGetVaultCount() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // Initially no vaults + uint256 countBefore = timelock.getVaultCount(user1); + assertEq(countBefore, 0, "Should have 0 vaults initially"); + + // Create 3 vaults + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + timelock.deposit{value: 2 ether}(unlockTime); + timelock.deposit{value: 3 ether}(unlockTime); + vm.stopPrank(); + + // Check count + uint256 countAfter = timelock.getVaultCount(user1); + assertEq(countAfter, 3, "Should have 3 vaults"); + } + + // 🎯 ADDITIONAL TESTS FOR COMPLETE COVERAGE + + function testCannotWithdrawFromInactiveVault() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // Deposit and withdraw to make vault inactive + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + vm.warp(unlockTime + 1); + timelock.withdraw(0); + + // Try to withdraw again from inactive vault + vm.expectRevert("Vault is not active"); + timelock.withdraw(0); + vm.stopPrank(); + } + + function testCannotWithdrawFromZeroBalanceVault() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // Deposit and withdraw + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + vm.warp(unlockTime + 1); + timelock.withdraw(0); + + // Vault is now inactive with zero balance + (uint256 balance,, bool active,) = timelock.getVault(user1, 0); + assertEq(balance, 0, "Vault balance should be 0"); + assertFalse(active, "Vault should be inactive"); + vm.stopPrank(); + } + + function testVaultStateChangesAfterWithdraw() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 5 ether; + + // Deposit + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Check initial state + (uint256 balanceBefore, uint256 unlockBefore, bool activeBefore, bool isUnlockedBefore) = + timelock.getVault(user1, 0); + assertEq(balanceBefore, depositAmount, "Initial balance should match deposit"); + assertEq(unlockBefore, unlockTime, "Unlock time should match"); + assertTrue(activeBefore, "Vault should be active"); + assertFalse(isUnlockedBefore, "Vault should be locked initially"); + + // Move time and withdraw + vm.warp(unlockTime + 1); + vm.prank(user1); + timelock.withdraw(0); + + // Check state after withdrawal + (uint256 balanceAfter,, bool activeAfter,) = timelock.getVault(user1, 0); + assertEq(balanceAfter, 0, "Balance should be 0 after withdrawal"); + assertFalse(activeAfter, "Vault should be inactive after withdrawal"); + } + + function testContractBalanceChanges() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // Initial contract balance should be 0 + assertEq(address(timelock).balance, 0, "Contract should start with 0 balance"); + + // User1 deposits 1 ether + vm.prank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + assertEq(address(timelock).balance, 1 ether, "Contract should have 1 ether"); + + // User2 deposits 2 ether + vm.prank(user2); + timelock.deposit{value: 2 ether}(unlockTime); + assertEq(address(timelock).balance, 3 ether, "Contract should have 3 ether"); + + // User1 withdraws + vm.warp(unlockTime + 1); + vm.prank(user1); + timelock.withdraw(0); + assertEq(address(timelock).balance, 2 ether, "Contract should have 2 ether after user1 withdrawal"); + + // User2 withdraws + vm.prank(user2); + timelock.withdraw(0); + assertEq(address(timelock).balance, 0, "Contract should have 0 ether after all withdrawals"); + } + + function testInitialUserBalances() public view { + // Check that users were given correct initial balances in setUp + assertEq(user1.balance, 1000 ether, "User1 should have 1000 ether initially"); + assertEq(user2.balance, 1000 ether, "User2 should have 1000 ether initially"); + } + + function testUserBalanceDecreasesOnDeposit() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 10 ether; + + // Check state BEFORE deposit + uint256 initialBalance = user1.balance; + assertEq(initialBalance, 1000 ether, "User should start with 1000 ether"); + + // Perform action: deposit + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Check state AFTER deposit + uint256 finalBalance = user1.balance; + + // Assert state changes + assertEq(initialBalance - finalBalance, depositAmount, "User balance should decrease by deposit amount"); + assertEq(finalBalance, 990 ether, "User should have 990 ether after depositing 10"); + } + + function testUserBalanceIncreasesOnWithdraw() public { + uint256 unlockTime = block.timestamp + 1 hours; + uint256 depositAmount = 10 ether; + + // Setup: Deposit + vm.prank(user1); + timelock.deposit{value: depositAmount}(unlockTime); + + // Check state BEFORE withdraw + uint256 balanceAfterDeposit = user1.balance; + assertEq(balanceAfterDeposit, 990 ether, "User should have 990 ether after deposit"); + + // Move time forward + vm.warp(unlockTime + 1); + + // Perform action: Withdraw + vm.prank(user1); + timelock.withdraw(0); + + // Check state AFTER withdraw + uint256 balanceAfterWithdraw = user1.balance; + + // Assert state changes + assertEq(balanceAfterWithdraw, 1000 ether, "User should have 1000 ether after withdrawal"); + assertEq(balanceAfterWithdraw - balanceAfterDeposit, depositAmount, "Balance increase should equal deposit amount"); + } + + function testMultipleDepositsStateChanges() public { + uint256 unlockTime = block.timestamp + 1 hours; + + // Check state BEFORE deposits + uint256 initialBalance = user1.balance; + uint256 initialContractBalance = address(timelock).balance; + uint256 initialVaultCount = timelock.getVaultCount(user1); + uint256 initialTotalBalance = timelock.getTotalBalance(user1); + + assertEq(initialBalance, 1000 ether, "User should start with 1000 ether"); + assertEq(initialContractBalance, 0, "Contract should start with 0 balance"); + assertEq(initialVaultCount, 0, "User should have 0 vaults initially"); + assertEq(initialTotalBalance, 0, "Total balance should be 0 initially"); + + // Perform action: Make 3 deposits + vm.startPrank(user1); + timelock.deposit{value: 1 ether}(unlockTime); + timelock.deposit{value: 2 ether}(unlockTime); + timelock.deposit{value: 3 ether}(unlockTime); + vm.stopPrank(); + + // Check state AFTER deposits + uint256 finalBalance = user1.balance; + uint256 finalContractBalance = address(timelock).balance; + uint256 finalVaultCount = timelock.getVaultCount(user1); + uint256 finalTotalBalance = timelock.getTotalBalance(user1); + + // Assert state changes + assertEq(initialBalance - finalBalance, 6 ether, "User balance should decrease by 6 ether"); + assertEq(finalContractBalance - initialContractBalance, 6 ether, "Contract balance should increase by 6 ether"); + assertEq(finalBalance, 994 ether, "User should have 994 ether"); + assertEq(finalContractBalance, 6 ether, "Contract should have 6 ether"); + assertEq(finalVaultCount, 3, "Should have 3 vaults"); + assertEq(finalTotalBalance, 6 ether, "Total balance should be 6 ether"); + } + + function testWithdrawAllStateChanges() public { + uint256 unlockTime1 = block.timestamp + 1 hours; + uint256 unlockTime2 = block.timestamp + 2 hours; + + // Setup: Create 2 vaults + vm.startPrank(user1); + timelock.deposit{value: 3 ether}(unlockTime1); + timelock.deposit{value: 5 ether}(unlockTime2); + vm.stopPrank(); + + // Unlock both + vm.warp(unlockTime2 + 1); + + // Check state BEFORE withdrawAll + uint256 balanceBeforeWithdraw = user1.balance; + uint256 contractBalanceBefore = address(timelock).balance; + uint256 totalBalanceBefore = timelock.getTotalBalance(user1); + uint256 unlockedBalanceBefore = timelock.getUnlockedBalance(user1); + (,, bool active1Before,) = timelock.getVault(user1, 0); + (,, bool active2Before,) = timelock.getVault(user1, 1); + + assertEq(balanceBeforeWithdraw, 992 ether, "User should have 992 ether after deposits"); + assertEq(contractBalanceBefore, 8 ether, "Contract should have 8 ether"); + assertEq(totalBalanceBefore, 8 ether, "Total balance should be 8 ether"); + assertEq(unlockedBalanceBefore, 8 ether, "All funds should be unlocked"); + assertTrue(active1Before, "Vault 1 should be active"); + assertTrue(active2Before, "Vault 2 should be active"); + + // Perform action: withdraw all + vm.prank(user1); + uint256 totalWithdrawn = timelock.withdrawAll(); + + // Check state AFTER withdrawAll + uint256 balanceAfterWithdraw = user1.balance; + uint256 contractBalanceAfter = address(timelock).balance; + uint256 totalBalanceAfter = timelock.getTotalBalance(user1); + (,, bool active1After,) = timelock.getVault(user1, 0); + (,, bool active2After,) = timelock.getVault(user1, 1); + + // Assert state changes + assertEq(totalWithdrawn, 8 ether, "Should withdraw 8 ether total"); + assertEq(balanceAfterWithdraw - balanceBeforeWithdraw, 8 ether, "User should receive 8 ether"); + assertEq(contractBalanceBefore - contractBalanceAfter, 8 ether, "Contract should lose 8 ether"); + assertEq(balanceAfterWithdraw, 1000 ether, "User should have 1000 ether back"); + assertEq(contractBalanceAfter, 0, "Contract should have 0 ether"); + assertEq(totalBalanceAfter, 0, "Total balance should be 0"); + assertFalse(active1After, "Vault 1 should be inactive"); + assertFalse(active2After, "Vault 2 should be inactive"); + } + + function testGetVaultRevertsForInvalidId() public { + // Try to get a vault that doesn't exist + vm.expectRevert("Invalid vault ID"); + timelock.getVault(user1, 0); + + // Create one vault + vm.prank(user1); + timelock.deposit{value: 1 ether}(block.timestamp + 1 hours); + + // Try to get vault ID 1 (doesn't exist) + vm.expectRevert("Invalid vault ID"); + timelock.getVault(user1, 1); + } +}