From c6eba3e80d86df77ebb0c4122491116f5febc995 Mon Sep 17 00:00:00 2001 From: Henry Peters Date: Fri, 6 Mar 2026 14:37:38 +0100 Subject: [PATCH 1/2] feat: implement timelockv1 --- .gitignore | 1 + .../timelock/.github/workflows/test.yml | 38 ++ assignments/foundry-test/timelock/.gitignore | 14 + assignments/foundry-test/timelock/.gitmodules | 3 + assignments/foundry-test/timelock/README.md | 66 +++ .../foundry-test/timelock/foundry.lock | 8 + .../foundry-test/timelock/foundry.toml | 6 + .../foundry-test/timelock/remappings.txt | 5 + .../timelock/script/Counter.s.sol | 19 + .../foundry-test/timelock/src/Counter.sol | 14 + .../foundry-test/timelock/src/IERC20.sol | 14 + .../foundry-test/timelock/src/TimelockV2.sol | 182 +++++++ .../foundry-test/timelock/src/token.sol | 17 + .../foundry-test/timelock/test/Counter.t.sol | 24 + .../timelock/test/TimelockV2.t.sol | 460 ++++++++++++++++++ 15 files changed, 871 insertions(+) create mode 100644 assignments/foundry-test/timelock/.github/workflows/test.yml create mode 100644 assignments/foundry-test/timelock/.gitignore create mode 100644 assignments/foundry-test/timelock/.gitmodules create mode 100644 assignments/foundry-test/timelock/README.md create mode 100644 assignments/foundry-test/timelock/foundry.lock create mode 100644 assignments/foundry-test/timelock/foundry.toml create mode 100644 assignments/foundry-test/timelock/remappings.txt create mode 100644 assignments/foundry-test/timelock/script/Counter.s.sol create mode 100644 assignments/foundry-test/timelock/src/Counter.sol create mode 100644 assignments/foundry-test/timelock/src/IERC20.sol create mode 100644 assignments/foundry-test/timelock/src/TimelockV2.sol create mode 100644 assignments/foundry-test/timelock/src/token.sol create mode 100644 assignments/foundry-test/timelock/test/Counter.t.sol create mode 100644 assignments/foundry-test/timelock/test/TimelockV2.t.sol diff --git a/.gitignore b/.gitignore index 60f73e90..0c8e95b5 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ node_modules **/target artifacts/ cache/ +lib/ typechain-types/ # Logs diff --git a/assignments/foundry-test/timelock/.github/workflows/test.yml b/assignments/foundry-test/timelock/.github/workflows/test.yml new file mode 100644 index 00000000..b79c8d4f --- /dev/null +++ b/assignments/foundry-test/timelock/.github/workflows/test.yml @@ -0,0 +1,38 @@ +name: CI + +permissions: {} + +on: + push: + pull_request: + workflow_dispatch: + +env: + FOUNDRY_PROFILE: ci + +jobs: + check: + name: Foundry project + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - uses: actions/checkout@v5 + with: + persist-credentials: false + submodules: recursive + + - name: Install Foundry + uses: foundry-rs/foundry-toolchain@v1 + + - name: Show Forge version + run: forge --version + + - name: Run Forge fmt + run: forge fmt --check + + - name: Run Forge build + run: forge build --sizes + + - name: Run Forge tests + run: forge test -vvv diff --git a/assignments/foundry-test/timelock/.gitignore b/assignments/foundry-test/timelock/.gitignore new file mode 100644 index 00000000..85198aaa --- /dev/null +++ b/assignments/foundry-test/timelock/.gitignore @@ -0,0 +1,14 @@ +# Compiler files +cache/ +out/ + +# Ignores development broadcast logs +!/broadcast +/broadcast/*/31337/ +/broadcast/**/dry-run/ + +# Docs +docs/ + +# Dotenv file +.env diff --git a/assignments/foundry-test/timelock/.gitmodules b/assignments/foundry-test/timelock/.gitmodules new file mode 100644 index 00000000..888d42dc --- /dev/null +++ b/assignments/foundry-test/timelock/.gitmodules @@ -0,0 +1,3 @@ +[submodule "lib/forge-std"] + path = lib/forge-std + url = https://github.com/foundry-rs/forge-std diff --git a/assignments/foundry-test/timelock/README.md b/assignments/foundry-test/timelock/README.md new file mode 100644 index 00000000..8817d6ab --- /dev/null +++ b/assignments/foundry-test/timelock/README.md @@ -0,0 +1,66 @@ +## Foundry + +**Foundry is a blazing fast, portable and modular toolkit for Ethereum application development written in Rust.** + +Foundry consists of: + +- **Forge**: Ethereum testing framework (like Truffle, Hardhat and DappTools). +- **Cast**: Swiss army knife for interacting with EVM smart contracts, sending transactions and getting chain data. +- **Anvil**: Local Ethereum node, akin to Ganache, Hardhat Network. +- **Chisel**: Fast, utilitarian, and verbose solidity REPL. + +## Documentation + +https://book.getfoundry.sh/ + +## Usage + +### Build + +```shell +$ forge build +``` + +### Test + +```shell +$ forge test +``` + +### Format + +```shell +$ forge fmt +``` + +### Gas Snapshots + +```shell +$ forge snapshot +``` + +### Anvil + +```shell +$ anvil +``` + +### Deploy + +```shell +$ forge script script/Counter.s.sol:CounterScript --rpc-url --private-key +``` + +### Cast + +```shell +$ cast +``` + +### Help + +```shell +$ forge --help +$ anvil --help +$ cast --help +``` diff --git a/assignments/foundry-test/timelock/foundry.lock b/assignments/foundry-test/timelock/foundry.lock new file mode 100644 index 00000000..bc06b89b --- /dev/null +++ b/assignments/foundry-test/timelock/foundry.lock @@ -0,0 +1,8 @@ +{ + "lib/forge-std": { + "tag": { + "name": "v1.15.0", + "rev": "0844d7e1fc5e60d77b68e469bff60265f236c398" + } + } +} \ No newline at end of file diff --git a/assignments/foundry-test/timelock/foundry.toml b/assignments/foundry-test/timelock/foundry.toml new file mode 100644 index 00000000..25b918f9 --- /dev/null +++ b/assignments/foundry-test/timelock/foundry.toml @@ -0,0 +1,6 @@ +[profile.default] +src = "src" +out = "out" +libs = ["lib"] + +# See more config options https://github.com/foundry-rs/foundry/blob/master/crates/config/README.md#all-options diff --git a/assignments/foundry-test/timelock/remappings.txt b/assignments/foundry-test/timelock/remappings.txt new file mode 100644 index 00000000..918ed315 --- /dev/null +++ b/assignments/foundry-test/timelock/remappings.txt @@ -0,0 +1,5 @@ +@openzeppelin/contracts/=lib/openzeppelin-contracts/contracts/ +erc4626-tests/=lib/openzeppelin-contracts/lib/erc4626-tests/ +forge-std/=lib/forge-std/src/ +halmos-cheatcodes/=lib/openzeppelin-contracts/lib/halmos-cheatcodes/src/ +openzeppelin-contracts/=lib/openzeppelin-contracts/ diff --git a/assignments/foundry-test/timelock/script/Counter.s.sol b/assignments/foundry-test/timelock/script/Counter.s.sol new file mode 100644 index 00000000..f01d69c3 --- /dev/null +++ b/assignments/foundry-test/timelock/script/Counter.s.sol @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import {Script} from "forge-std/Script.sol"; +import {Counter} from "../src/Counter.sol"; + +contract CounterScript is Script { + Counter public counter; + + function setUp() public {} + + function run() public { + vm.startBroadcast(); + + counter = new Counter(); + + vm.stopBroadcast(); + } +} diff --git a/assignments/foundry-test/timelock/src/Counter.sol b/assignments/foundry-test/timelock/src/Counter.sol new file mode 100644 index 00000000..aded7997 --- /dev/null +++ b/assignments/foundry-test/timelock/src/Counter.sol @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +contract Counter { + uint256 public number; + + function setNumber(uint256 newNumber) public { + number = newNumber; + } + + function increment() public { + number++; + } +} diff --git a/assignments/foundry-test/timelock/src/IERC20.sol b/assignments/foundry-test/timelock/src/IERC20.sol new file mode 100644 index 00000000..9df9b1c3 --- /dev/null +++ b/assignments/foundry-test/timelock/src/IERC20.sol @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; + +interface IErc20 is IERC20 { + function decimals() external view returns (uint8); + function totalSupply() external view returns (uint256); + function balanceOf(address _owner) external view returns (uint256 balance); + function transferFrom(address _from, address _to, uint256 _value) external returns (bool success); + function mint(address to, uint256 amount) external; + function burn(uint256 amount) external; + function burnFrom(address account, uint256 amount) external; +} diff --git a/assignments/foundry-test/timelock/src/TimelockV2.sol b/assignments/foundry-test/timelock/src/TimelockV2.sol new file mode 100644 index 00000000..7c717f18 --- /dev/null +++ b/assignments/foundry-test/timelock/src/TimelockV2.sol @@ -0,0 +1,182 @@ +//SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.28; + +import {IErc20} from "./IERC20.sol"; + +contract TimeLockV2 { + IErc20 public immutable henToken; + address public owner; + + struct Vault { + uint256 balance; + uint256 tokenBalance; + uint256 unlockTime; + bool active; + } + + constructor(address _token) { + henToken = IErc20(_token); + owner = msg.sender; + } + + modifier onlyOwner() { + require(msg.sender == owner, "Not owner"); + _; + } + + mapping(address => Vault[]) private vaults; + + event Deposited(address indexed user, uint256 vaultId, uint256 amount, uint256 unlockTime); + event Withdrawn(address indexed user, uint256 vaultId, uint256 amount); + + // Internal function + function _depositRatio(uint256 _totalDeposit) internal pure returns (uint256 _tokenAmount) { + _tokenAmount = _totalDeposit * 10; // Token Ratio + } + + function deposit(uint256 _unlockTime) external payable returns (uint256) { + require(msg.value > 0, "Deposit must be greater than zero"); + require(_unlockTime > block.timestamp, "Unlock time must be in the future"); + + uint256 tokenBal = _depositRatio(msg.value); + require(henToken.transfer(msg.sender, tokenBal), "Token transfer failed"); + + // Create new vault + vaults[msg.sender].push( + Vault({balance: msg.value, unlockTime: _unlockTime, tokenBalance: tokenBal, active: true}) + ); + + uint256 vaultId = vaults[msg.sender].length - 1; + emit Deposited(msg.sender, vaultId, msg.value, _unlockTime); + + return vaultId; + } + + function withdraw(uint256 _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"); + + uint256 amount = userVault.balance; + uint256 tokenAmount = userVault.tokenBalance; + + // Mark vault as inactive and clear balance + userVault.balance = 0; + userVault.tokenBalance = 0; + userVault.active = false; + + require(henToken.transferFrom(msg.sender, address(this), tokenAmount), "Token transfer failed"); + + (bool success,) = payable(msg.sender).call{value: amount}(""); + require(success, "Transfer failed"); + + emit Withdrawn(msg.sender, _vaultId, amount); + } + + function withdrawAll() external returns (uint256) { + uint256 totalWithdrawn = 0; + Vault[] storage userVaults = vaults[msg.sender]; + + for (uint256 i = 0; i < userVaults.length; i++) { + if (userVaults[i].active && userVaults[i].balance > 0 && block.timestamp >= userVaults[i].unlockTime) { + uint256 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 emergencyWithdraw() external onlyOwner returns (uint256 amount) { + amount = address(this).balance; + require(amount > 0, "No funds available"); + + (bool success,) = payable(owner).call{value: amount}(""); + if (!success) revert(); + } + + function getVaultCount(address _user) external view returns (uint256) { + return vaults[_user].length; + } + + function getVault(address _user, uint256 _vaultId) + external + view + returns (uint256 balance, uint256 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 (uint256[] memory activeVaults, uint256[] memory balances, uint256[] memory unlockTimes) + { + Vault[] storage userVaults = vaults[_user]; + + // Count active vaults + uint256 activeCount = 0; + for (uint256 i = 0; i < userVaults.length; i++) { + if (userVaults[i].active && userVaults[i].balance > 0) { + activeCount++; + } + } + + // Create arrays + activeVaults = new uint256[](activeCount); + balances = new uint256[](activeCount); + unlockTimes = new uint256[](activeCount); + + // Populate arrays + uint256 index = 0; + for (uint256 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 (uint256 total) { + Vault[] storage userVaults = vaults[_user]; + for (uint256 i = 0; i < userVaults.length; i++) { + if (userVaults[i].active) { + total += userVaults[i].balance; + } + } + return total; + } + + function getUnlockedBalance(address _user) external view returns (uint256 unlocked) { + Vault[] storage userVaults = vaults[_user]; + for (uint256 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/timelock/src/token.sol b/assignments/foundry-test/timelock/src/token.sol new file mode 100644 index 00000000..6f5d88f5 --- /dev/null +++ b/assignments/foundry-test/timelock/src/token.sol @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: MIT +// Compatible with OpenZeppelin Contracts ^5.6.0 +pragma solidity ^0.8.27; + +import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol"; +import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; +import {ERC20Burnable} from "@openzeppelin/contracts/token/ERC20/extensions/ERC20Burnable.sol"; + +contract HENTU is ERC20, ERC20Burnable, Ownable { + constructor(address recipient, address initialOwner) ERC20("HENTU", "HTU") Ownable(initialOwner) { + _mint(recipient, 1000 * 10 ** decimals()); + } + + function mint(address to, uint256 amount) public onlyOwner { + _mint(to, amount); + } +} diff --git a/assignments/foundry-test/timelock/test/Counter.t.sol b/assignments/foundry-test/timelock/test/Counter.t.sol new file mode 100644 index 00000000..48319108 --- /dev/null +++ b/assignments/foundry-test/timelock/test/Counter.t.sol @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import {Test} from "forge-std/Test.sol"; +import {Counter} from "../src/Counter.sol"; + +contract CounterTest is Test { + Counter public counter; + + function setUp() public { + counter = new Counter(); + counter.setNumber(0); + } + + function test_Increment() public { + counter.increment(); + assertEq(counter.number(), 1); + } + + function testFuzz_SetNumber(uint256 x) public { + counter.setNumber(x); + assertEq(counter.number(), x); + } +} diff --git a/assignments/foundry-test/timelock/test/TimelockV2.t.sol b/assignments/foundry-test/timelock/test/TimelockV2.t.sol new file mode 100644 index 00000000..ca365747 --- /dev/null +++ b/assignments/foundry-test/timelock/test/TimelockV2.t.sol @@ -0,0 +1,460 @@ +//SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {console} from "forge-std/console.sol"; +import {TimeLockV2} from "../src/TimelockV2.sol"; +import {HENTU} from "../src/token.sol"; + +contract TimeLockV2Test is Test { + TimeLockV2 public timelock; + HENTU public henToken; + uint256 amountMinted; + + address public addr1 = makeAddr("addr1"); + address public addr2 = makeAddr("addr2"); + address public owner = makeAddr("owner"); + + uint256 public constant DEPOSIT_AMOUNT = 1 ether; + uint256 public constant TOKEN_RATIO = 10; + + // Events from TimeLockV2 + event Deposited(address indexed user, uint256 vaultId, uint256 amount, uint256 unlockTime); + event Withdrawn(address indexed user, uint256 vaultId, uint256 amount); + + function setUp() public { + vm.startPrank(owner); + henToken = new HENTU(address(this), address(this)); + timelock = new TimeLockV2(address(henToken)); + vm.stopPrank(); + + amountMinted = 1000000 * 10 ** henToken.decimals(); + henToken.mint(address(timelock), amountMinted); + } + + function test_emergencyWithdraw() public { + uint256 unlockTime = block.timestamp + 1 days; + uint256 ownerBalanceBefore = owner.balance; + + vm.deal(addr2, 1 ether); + vm.prank(addr2); + timelock.deposit{value: 1 ether}(unlockTime); + + vm.prank(owner); + uint256 withdrawnAmount = timelock.emergencyWithdraw(); + + assertEq(withdrawnAmount, 1 ether); + assertEq(address(timelock).balance, 0); + assertEq(owner.balance, ownerBalanceBefore + 1 ether); + } + + function testDeposit() public { + vm.deal(addr1, 10 ether); + + uint256 depositAmount = 1 ether; + uint256 tokenRatio = 10; + + uint256 unlockTime = block.timestamp + 1 days; + uint256 expectedTokens = depositAmount * tokenRatio; + + vm.startPrank(addr1); + + ////// + uint256 addresTokenBalanceBefore = henToken.balanceOf(addr1); + assertEq(addresTokenBalanceBefore, 0); + + uint256 addressETHBalanceBefore = addr1.balance; + assertEq(addressETHBalanceBefore, 10 ether); + + uint256 contractTokenBalanceBefore = henToken.balanceOf(address(timelock)); + assertEq(contractTokenBalanceBefore, amountMinted); + + uint256 contractETHBalanceBefore = address(timelock).balance; + assertEq(contractETHBalanceBefore, 0); + /////// + + timelock.deposit{value: depositAmount}(unlockTime); + + uint256 vaultId = timelock.getVaultCount(addr1); + // assert that a vault has been created after deposit + assertEq(vaultId, 1); + + ////// + uint256 addresTokenBalanceAfter = henToken.balanceOf(addr1); + assertEq(addresTokenBalanceAfter, expectedTokens); + + uint256 addressETHBalanceAfter = addr1.balance; + assertEq(addressETHBalanceAfter, 9 ether); + + uint256 contractTokenBalanceAfter = henToken.balanceOf(address(timelock)); + assertEq(contractTokenBalanceAfter, amountMinted - expectedTokens); + + uint256 contractETHBalanceAfter = address(timelock).balance; + assertEq(contractETHBalanceAfter, 1 ether); + /////// + + assertEq(timelock.getVaultCount(addr1), 1); + + vm.stopPrank(); + } + + // Test deposit with zero value fails + function testDepositZeroValueFails() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + vm.expectRevert("Deposit must be greater than zero"); + timelock.deposit{value: 0}(unlockTime); + + vm.stopPrank(); + } + + // Test deposit with past unlock time fails + function testDepositPastUnlockTimeFails() public { + vm.deal(addr1, 10 ether); + uint256 depositAmount = 1 ether; + + uint256 unlockTime = block.timestamp - 1; // Past time + + vm.startPrank(addr1); + + vm.expectRevert("Unlock time must be in the future"); + timelock.deposit{value: depositAmount}(unlockTime); + + vm.stopPrank(); + } + + // Test successful withdrawal + function testWithdraw() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 1 days; + uint256 expectedTokens = DEPOSIT_AMOUNT * TOKEN_RATIO; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + uint256 vaultId = timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + uint256 userBalanceBefore = addr1.balance; + uint256 userTokenBalanceBefore = henToken.balanceOf(addr1); + + vm.warp(unlockTime + 1); + + timelock.withdraw(vaultId); + + // Check ETH was returned + assertEq(addr1.balance, userBalanceBefore + DEPOSIT_AMOUNT); + + // Check tokens were burned (transferred back) + assertEq(henToken.balanceOf(addr1), userTokenBalanceBefore - expectedTokens); + + vm.stopPrank(); + } + + // Test withdrawal before unlock time fails + function testWithdrawBeforeUnlockFails() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 7 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + uint256 vaultId = timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + vm.expectRevert("Funds are still locked"); + timelock.withdraw(vaultId); + + vm.stopPrank(); + } + + // Test withdrawal with invalid vault ID fails + function testWithdrawInvalidVaultIdFails() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + vm.expectRevert("Invalid vault ID"); + timelock.withdraw(999); // Invalid vault ID + + vm.stopPrank(); + } + + // Test withdrawAll + function testWithdrawAll() public { + vm.deal(addr1, 100 ether); + + uint256 unlockTime1 = block.timestamp + 1 days; + uint256 unlockTime2 = block.timestamp + 2 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + // Create multiple vaults + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime1); + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime2); + + // Warp to after first unlock time + vm.warp(unlockTime1 + 1); + + uint256 balanceBefore = addr1.balance; + + // Withdraw all unlocked + uint256 withdrawn = timelock.withdrawAll(); + + // Only first vault should be unlocked + assertEq(withdrawn, DEPOSIT_AMOUNT); + assertEq(addr1.balance, balanceBefore + DEPOSIT_AMOUNT); + + vm.stopPrank(); + } + + // Test withdrawAll with no unlocked funds fails + function testWithdrawAllNoUnlockedFails() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 7 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + vm.expectRevert("No unlocked funds available"); + timelock.withdrawAll(); + + vm.stopPrank(); + } + + // Test getVault function + function testGetVault() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + uint256 vaultId = timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + (uint256 balance, uint256 unlockTimeRet, bool active, bool isUnlocked) = timelock.getVault(addr1, vaultId); + + assertEq(balance, DEPOSIT_AMOUNT); + assertEq(unlockTimeRet, unlockTime); + assertTrue(active); + assertFalse(isUnlocked); // Not unlocked yet + + // Warp past unlock time + vm.warp(unlockTime + 1); + + (,,, isUnlocked) = timelock.getVault(addr1, vaultId); + assertTrue(isUnlocked); + + vm.stopPrank(); + } + + // Test getVaultCount + function testGetVaultCount() public { + vm.deal(addr1, 100 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + assertEq(timelock.getVaultCount(addr1), 0); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + assertEq(timelock.getVaultCount(addr1), 1); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + assertEq(timelock.getVaultCount(addr1), 2); + + vm.stopPrank(); + } + + // Test getTotalBalance + function testGetTotalBalance() public { + vm.deal(addr1, 100 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + assertEq(timelock.getTotalBalance(addr1), DEPOSIT_AMOUNT * 2); + + vm.stopPrank(); + } + + // Test getUnlockedBalance + function testGetUnlockedBalance() public { + vm.deal(addr1, 100 ether); + + uint256 unlockTime1 = block.timestamp + 1 days; + uint256 unlockTime2 = block.timestamp + 2 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime1); + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime2); + + // Initially no unlocked balance + assertEq(timelock.getUnlockedBalance(addr1), 0); + + // Warp past first unlock time + vm.warp(unlockTime1 + 1); + + // Should have one vault unlocked + assertEq(timelock.getUnlockedBalance(addr1), DEPOSIT_AMOUNT); + + // Warp past second unlock time + vm.warp(unlockTime2 + 1); + + // Should have both vaults unlocked + assertEq(timelock.getUnlockedBalance(addr1), DEPOSIT_AMOUNT * 2); + + vm.stopPrank(); + } + + // Test getActiveVaults + function testGetActiveVaults() public { + vm.deal(addr1, 100 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + (uint256[] memory activeVaults, uint256[] memory balances, uint256[] memory unlockTimes) = + timelock.getActiveVaults(addr1); + + assertEq(activeVaults.length, 2); + assertEq(balances.length, 2); + assertEq(unlockTimes.length, 2); + assertEq(activeVaults[0], 0); + assertEq(activeVaults[1], 1); + assertEq(balances[0], DEPOSIT_AMOUNT); + assertEq(balances[1], DEPOSIT_AMOUNT); + + vm.stopPrank(); + } + + // Test getAllVaults + function testGetAllVaults() public { + vm.deal(addr1, 100 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + TimeLockV2.Vault[] memory allVaults = timelock.getAllVaults(addr1); + + assertEq(allVaults.length, 2); + assertEq(allVaults[0].balance, DEPOSIT_AMOUNT); + assertEq(allVaults[1].balance, DEPOSIT_AMOUNT); + assertTrue(allVaults[0].active); + assertTrue(allVaults[1].active); + + vm.stopPrank(); + } + + // Test multiple users have separate vaults + function testMultipleUsersSeparateVaults() public { + vm.deal(addr1, 10 ether); + vm.deal(addr2, 10 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + // addr1 deposits + vm.startPrank(addr1); + henToken.approve(address(timelock), type(uint256).max); + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + vm.stopPrank(); + + // addr2 deposits + vm.startPrank(addr2); + henToken.approve(address(timelock), type(uint256).max); + timelock.deposit{value: DEPOSIT_AMOUNT * 2}(unlockTime); + vm.stopPrank(); + + // Check vault counts + assertEq(timelock.getVaultCount(addr1), 1); + assertEq(timelock.getVaultCount(addr2), 1); + + // Check balances + assertEq(timelock.getTotalBalance(addr1), DEPOSIT_AMOUNT); + assertEq(timelock.getTotalBalance(addr2), DEPOSIT_AMOUNT * 2); + } + + // Test deposit event + function testDepositEvent() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + // vm.expectEmit(true, true, true, true); + emit Deposited(addr1, 0, DEPOSIT_AMOUNT, unlockTime); + + timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + vm.stopPrank(); + } + + // Test withdraw event + function testWithdrawEvent() public { + vm.deal(addr1, 10 ether); + + uint256 unlockTime = block.timestamp + 1 days; + + vm.startPrank(addr1); + + henToken.approve(address(timelock), type(uint256).max); + + uint256 vaultId = timelock.deposit{value: DEPOSIT_AMOUNT}(unlockTime); + + vm.warp(unlockTime + 1); + + vm.expectEmit(true, true, true, true); + emit Withdrawn(addr1, vaultId, DEPOSIT_AMOUNT); + + timelock.withdraw(vaultId); + + vm.stopPrank(); + } +} From f03192602930b40364d0571fb3ed741611667546 Mon Sep 17 00:00:00 2001 From: Henry Peters Date: Fri, 6 Mar 2026 14:46:45 +0100 Subject: [PATCH 2/2] chore: add out/ to .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 0c8e95b5..c3b5777d 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ node_modules artifacts/ cache/ lib/ +out/ typechain-types/ # Logs