smom-dbis-138/test/flash/SimpleERC3156FlashVault.t.sol

// 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);
    }
}