// SPDX-License-Identifier: MIT pragma solidity ^0.8.20; import {Test} from "forge-std/Test.sol"; import {TreasuryVault} from "../../contracts/treasury/TreasuryVault.sol"; contract MockERC20 { mapping(address => uint256) public balanceOf; function mint(address to, uint256 amount) external { balanceOf[to] += amount; } function transfer(address to, uint256 amount) external returns (bool) { balanceOf[msg.sender] -= amount; balanceOf[to] += amount; return true; } function transferFrom(address from, address to, uint256 amount) external returns (bool) { balanceOf[from] -= amount; balanceOf[to] += amount; return true; } function approve(address, uint256) external pure returns (bool) { return true; } } contract TreasuryVaultTest is Test { TreasuryVault vault; MockERC20 token; address admin = address(1); address module1 = address(2); address module2 = address(3); function setUp() public { vault = new TreasuryVault(admin); token = new MockERC20(); token.mint(address(vault), 1000000e18); vm.prank(admin); vault.setToken(address(token), true); vm.prank(admin); vault.setModule(module1, true); vm.prank(admin); vault.setModule(module2, true); } function test_requestTransfer_success() public { token.mint(address(vault), 100e18); vm.prank(module1); vault.requestTransfer(address(token), 50e18, module2); assertEq(token.balanceOf(module2), 50e18); } function test_requestTransfer_revert_tokenNotApproved() public { MockERC20 other = new MockERC20(); other.mint(address(vault), 100e18); vm.prank(module1); vm.expectRevert(TreasuryVault.TokenNotApproved.selector); vault.requestTransfer(address(other), 50e18, module2); } }