// SPDX-License-Identifier: MIT pragma solidity ^0.8.20; import {Test} from "forge-std/Test.sol"; import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; import {IERC3156FlashBorrower} from "@openzeppelin/contracts/interfaces/IERC3156FlashBorrower.sol"; import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; import {SimpleERC3156FlashVault} from "../../contracts/flash/SimpleERC3156FlashVault.sol"; contract MockBorrower is IERC3156FlashBorrower { bytes32 private constant _RETURN_VALUE = keccak256("ERC3156FlashBorrower.onFlashLoan"); function onFlashLoan(address, address token, uint256 amount, uint256 fee, bytes calldata) external override returns (bytes32) { IERC20(token).transfer(msg.sender, amount + fee); return _RETURN_VALUE; } } contract MockERC20Mint is ERC20 { constructor() ERC20("Mock", "MCK") {} function mint(address to, uint256 v) external { _mint(to, v); } } contract SimpleERC3156FlashVaultTest is Test { SimpleERC3156FlashVault internal vault; MockERC20Mint internal token; MockBorrower internal borrower; address internal owner = address(0xA11); address internal user = address(0xB22); function setUp() public { vm.startPrank(owner); vault = new SimpleERC3156FlashVault(owner, 5); // 0.05% token = new MockERC20Mint(); token.mint(address(vault), 1_000_000e18); vault.setTokenSupported(address(token), true); vm.stopPrank(); borrower = new MockBorrower(); token.mint(address(borrower), 100e18); vm.prank(address(borrower)); token.approve(address(vault), type(uint256).max); } function test_maxFlashLoan() public view { assertEq(vault.maxFlashLoan(address(token)), 1_000_000e18); } function test_flashFee() public view { assertEq(vault.flashFee(address(token), 100_000e18), (100_000e18 * 5) / 10_000); } function test_previewFlashFee_matches_flashFee() public view { uint256 a = 123_456e18; assertEq(vault.previewFlashFee(address(token), a), vault.flashFee(address(token), a)); } function test_flashLoan_happyPath() public { uint256 amount = 100_000e18; uint256 fee = vault.flashFee(address(token), amount); uint256 beforeBal = token.balanceOf(address(vault)); vm.prank(user); vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(token), amount, ""); assertEq(token.balanceOf(address(vault)), beforeBal + fee); assertEq(vault.totalFeesCollected(address(token)), fee); } function test_flashLoan_revert_unsupported() public { MockERC20Mint other = new MockERC20Mint(); vm.expectRevert(SimpleERC3156FlashVault.UnsupportedToken.selector); vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(other), 1, ""); } function test_setFeeBps_revert_above_max() public { vm.prank(owner); vm.expectRevert(SimpleERC3156FlashVault.FeeTooHigh.selector); vault.setFeeBps(1001); } function test_rescueTokens() public { vm.prank(owner); vault.rescueTokens(address(token), 10e18, owner); assertEq(token.balanceOf(owner), 10e18); } function test_borrowerAllowlist_revert_unapproved() public { vm.prank(owner); vault.setBorrowerAllowlistEnabled(true); vm.expectRevert(SimpleERC3156FlashVault.BorrowerNotApproved.selector); vm.prank(user); vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(token), 1e18, ""); } function test_borrowerAllowlist_allows_approved() public { vm.startPrank(owner); vault.setBorrowerAllowlistEnabled(true); vault.setBorrowerApproved(address(borrower), true); vm.stopPrank(); uint256 amount = 1000e18; uint256 fee = vault.flashFee(address(token), amount); uint256 beforeBal = token.balanceOf(address(vault)); vm.prank(user); vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(token), amount, ""); assertEq(token.balanceOf(address(vault)), beforeBal + fee); } }