smom-dbis-138/test/flash/CrossChainFlashRepayReceiver.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 {CrossChainFlashRepayReceiver} from "../../contracts/flash/CrossChainFlashRepayReceiver.sol";
import {IRouterClient} from "../../contracts/ccip/IRouterClient.sol";

contract MockToken is ERC20 {
    constructor() ERC20("X", "X") {}

    function mint(address to, uint256 v) external {
        _mint(to, v);
    }
}

contract CrossChainFlashRepayReceiverTest is Test {
    CrossChainFlashRepayReceiver internal recv;
    MockToken internal token;
    address internal router;

    function setUp() public {
        router = makeAddr("ccipRouter");
        recv = new CrossChainFlashRepayReceiver(router);
        token = new MockToken();
    }

    function _msg(address beneficiary, bytes32 obligation, uint256 amt, address tok)
        internal
        pure
        returns (IRouterClient.Any2EVMMessage memory m)
    {
        IRouterClient.TokenAmount[] memory amounts = new IRouterClient.TokenAmount[](1);
        amounts[0] = IRouterClient.TokenAmount({token: tok, amount: amt, amountType: IRouterClient.TokenAmountType.Fiat});
        m = IRouterClient.Any2EVMMessage({
            messageId: keccak256("mid"),
            sourceChainSelector: 138,
            sender: abi.encode(address(0x111)),
            data: abi.encode(beneficiary, obligation),
            tokenAmounts: amounts
        });
    }

    function test_ccipReceive_forwardsToRecipient() public {
        address beneficiary = address(0xB0B);
        bytes32 obligation = keccak256("obligation-1");
        uint256 amt = 777e18;
        token.mint(address(recv), amt);

        IRouterClient.Any2EVMMessage memory message = _msg(beneficiary, obligation, amt, address(token));

        vm.prank(router);
        recv.ccipReceive(message);

        assertEq(token.balanceOf(beneficiary), amt);
    }

    function test_ccipReceive_revert_notRouter() public {
        IRouterClient.Any2EVMMessage memory message = _msg(address(0x1), bytes32(0), 1, address(token));
        vm.expectRevert(CrossChainFlashRepayReceiver.OnlyRouter.selector);
        recv.ccipReceive(message);
    }
}