strategic/contracts/test/AtomicExecutorFlashLoan.t.sol

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

import {Test, console} from "forge-std/Test.sol";
import {AtomicExecutor} from "../AtomicExecutor.sol";

// Mock Aave Pool for flash loan testing
contract MockAavePool {
    AtomicExecutor public executor;
    address public asset;
    uint256 public amount;
    bool public callbackExecuted = false;

    function setExecutor(address _executor) external {
        executor = AtomicExecutor(_executor);
    }

    function flashLoanSimple(
        address receiverAddress,
        address _asset,
        uint256 _amount,
        bytes calldata params,
        uint16
    ) external {
        asset = _asset;
        amount = _amount;
        
        // Transfer asset to receiver (simulating flash loan)
        IERC20(_asset).transfer(receiverAddress, _amount);
        
        // Call executeOperation callback
        IFlashLoanSimpleReceiver(receiverAddress).executeOperation(
            _asset,
            _amount,
            _amount / 1000, // 0.1% premium
            params
        );
        
        // Require repayment
        uint256 repayment = _amount + (_amount / 1000);
        require(
            IERC20(_asset).balanceOf(receiverAddress) >= repayment,
            "Insufficient repayment"
        );
        
        // Transfer repayment back
        IERC20(_asset).transferFrom(receiverAddress, address(this), repayment);
        
        callbackExecuted = true;
    }
}

// Mock ERC20 for testing
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;
    }
}

interface IFlashLoanSimpleReceiver {
    function executeOperation(
        address asset,
        uint256 amount,
        uint256 premium,
        bytes calldata params
    ) external returns (bool);
}

interface IERC20 {
    function balanceOf(address) external view returns (uint256);
    function transfer(address, uint256) external returns (bool);
    function transferFrom(address, address, uint256) external returns (bool);
}

contract AtomicExecutorFlashLoanTest is Test {
    AtomicExecutor executor;
    MockAavePool pool;
    MockERC20 token;
    address owner = address(1);
    address user = address(2);

    function setUp() public {
        vm.prank(owner);
        executor = new AtomicExecutor(owner);
        
        pool = new MockAavePool();
        token = new MockERC20();
        
        // Mint tokens to pool
        token.mint(address(pool), 1000000e18);
        
        // Set executor in pool
        pool.setExecutor(address(executor));
        
        // Allow pool for flash loans
        vm.prank(owner);
        executor.setAllowedPool(address(pool), true);
        
        // Allow executor to receive tokens
        vm.prank(owner);
        executor.setAllowedTarget(address(token), true);
    }

    function testExecuteFlashLoan() public {
        uint256 loanAmount = 1000e18;
        
        // Encode callback operations (empty for this test)
        bytes memory params = abi.encode(new address[](0), new bytes[](0));
        
        vm.prank(user);
        executor.executeFlashLoan(
            address(pool),
            address(token),
            loanAmount,
            params
        );
        
        assertTrue(pool.callbackExecuted());
    }

    function testFlashLoanRepayment() public {
        uint256 loanAmount = 1000e18;
        uint256 premium = loanAmount / 1000; // 0.1%
        uint256 repayment = loanAmount + premium;
        
        // Mint tokens to executor for repayment
        token.mint(address(executor), repayment);
        
        bytes memory params = abi.encode(new address[](0), new bytes[](0));
        
        vm.prank(user);
        executor.executeFlashLoan(
            address(pool),
            address(token),
            loanAmount,
            params
        );
        
        // Check pool has repayment
        assertGe(token.balanceOf(address(pool)), repayment);
    }

    function testFlashLoanUnauthorizedPool() public {
        MockAavePool unauthorizedPool = new MockAavePool();
        unauthorizedPool.setExecutor(address(executor));
        
        bytes memory params = abi.encode(new address[](0), new bytes[](0));
        
        vm.prank(user);
        vm.expectRevert("Unauthorized pool");
        executor.executeFlashLoan(
            address(unauthorizedPool),
            address(token),
            1000e18,
            params
        );
    }

  function testFlashLoanUnauthorizedInitiator() public {
    address attacker = address(999);
    
    bytes memory params = abi.encode(new address[](0), new bytes[](0));
    
    // Try to call executeOperation directly (should fail)
    vm.prank(address(pool));
    vm.expectRevert("Unauthorized initiator");
    executor.executeOperation(
      address(token),
      1000e18,
      1e18,
      params
    );
  }

  function testFlashLoanWithMultipleOperations() public {
    uint256 loanAmount = 1000e18;
    
    // Encode multiple operations in callback
    address[] memory targets = new address[](2);
    bytes[] memory calldatas = new bytes[](2);
    
    targets[0] = address(token);
    calldatas[0] = abi.encodeWithSignature("transfer(address,uint256)", address(1), 500e18);
    
    targets[1] = address(token);
    calldatas[1] = abi.encodeWithSignature("transfer(address,uint256)", address(2), 500e18);
    
    bytes memory params = abi.encode(targets, calldatas);
    
    vm.prank(user);
    executor.executeFlashLoan(
      address(pool),
      address(token),
      loanAmount,
      params
    );
    
    assertTrue(pool.callbackExecuted());
  }

  function testFlashLoanRepaymentValidation() public {
    uint256 loanAmount = 1000e18;
    uint256 premium = loanAmount / 1000; // 0.1%
    uint256 requiredRepayment = loanAmount + premium;
    
    // Don't mint enough for repayment
    token.mint(address(executor), loanAmount); // Not enough!
    
    bytes memory params = abi.encode(new address[](0), new bytes[](0));
    
    vm.prank(user);
    // Should revert due to insufficient repayment
    vm.expectRevert();
    executor.executeFlashLoan(
      address(pool),
      address(token),
      loanAmount,
      params
    );
  }
}