smom-dbis-138/test/WETH10.t.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.19;

import {Test, console} from "forge-std/Test.sol";
import {WETH10, IERC3156FlashBorrower} from "../contracts/tokens/WETH10.sol";

contract WETH10Test is Test {
    WETH10 public weth10;
    address public user = address(0x10); // Use address that can receive ETH (not precompile)
    address public recipient = address(2);
    
    function setUp() public {
        weth10 = new WETH10();
        vm.deal(user, 10 ether);
    }
    
    function testDeposit() public {
        vm.prank(user);
        weth10.deposit{value: 1 ether}();
        
        assertEq(weth10.balanceOf(user), 1 ether);
        assertEq(weth10.totalSupply(), 1 ether);
    }
    
    function testWithdraw() public {
        vm.deal(user, 10 ether);
        vm.prank(user);
        weth10.deposit{value: 1 ether}();
        
        // WETH10 contract now has 1 ether from the deposit
        uint256 balanceBefore = user.balance;
        vm.prank(user);
        weth10.withdraw(1 ether);
        
        assertEq(weth10.balanceOf(user), 0);
        assertEq(user.balance, balanceBefore + 1 ether);
    }
    
    function testTransfer() public {
        vm.prank(user);
        weth10.deposit{value: 1 ether}();
        
        vm.prank(user);
        weth10.transfer(recipient, 0.5 ether);
        
        assertEq(weth10.balanceOf(user), 0.5 ether);
        assertEq(weth10.balanceOf(recipient), 0.5 ether);
    }
    
    function testApprove() public {
        address spender = address(2);
        
        vm.prank(user);
        weth10.deposit{value: 1 ether}();
        
        vm.prank(user);
        weth10.approve(spender, 0.5 ether);
        
        assertEq(weth10.allowance(user, spender), 0.5 ether);
    }
    
    function testReceive() public {
        vm.prank(user);
        (bool success, ) = address(weth10).call{value: 1 ether}("");
        require(success, "Transfer failed");
        
        assertEq(weth10.balanceOf(user), 1 ether);
    }
    
    function testMaxFlashLoan() public {
        vm.prank(user);
        weth10.deposit{value: 5 ether}();
        
        uint256 maxLoan = weth10.maxFlashLoan(address(weth10));
        assertEq(maxLoan, 5 ether);
        
        uint256 maxLoanOther = weth10.maxFlashLoan(address(0));
        assertEq(maxLoanOther, 0);
    }
    
    function testFlashFee() public {
        uint256 fee = weth10.flashFee(address(weth10), 1 ether);
        assertEq(fee, 0); // WETH10 has no flash loan fee
    }
    
    function testFlashLoan() public {
        vm.deal(user, 10 ether);
        vm.prank(user);
        weth10.deposit{value: 10 ether}();
        
        FlashLoanReceiver receiver = new FlashLoanReceiver(weth10);
        
        uint256 amount = 5 ether;
        bytes memory data = "";
        
        vm.prank(user);
        bool success = weth10.flashLoan(receiver, address(weth10), amount, data);
        assertTrue(success);
    }
}

contract FlashLoanReceiver is IERC3156FlashBorrower {
    WETH10 public weth10;
    bytes32 public constant CALLBACK_SUCCESS = keccak256("ERC3156FlashBorrower.onFlashLoan");
    
    constructor(WETH10 _weth10) {
        weth10 = _weth10;
    }
    
    function onFlashLoan(
        address initiator,
        address token,
        uint256 amount,
        uint256 fee,
        bytes calldata data
    ) external override returns (bytes32) {
        require(msg.sender == address(weth10), "FlashLoanReceiver: invalid sender");
        require(token == address(weth10), "FlashLoanReceiver: invalid token");
        
        // Use the flash loan (e.g., arbitrage, liquidate, etc.)
        // For testing, we just repay immediately
        // The flash loan mints `amount` WETH to this contract
        // We need to have enough balance to repay `amount + fee`
        // Since fee is 0, we already have `amount` balance from the mint
        // No transfer needed - the balance is already sufficient
        
        return CALLBACK_SUCCESS;
    }
    
    function executeFlashLoan(uint256 amount) external {
        bytes memory data = "";
        weth10.flashLoan(this, address(weth10), amount, data);
    }
}