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

contract CreditMockToken is ERC20 {
    constructor() ERC20("Y", "Y") {}

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

contract CrossChainFlashVaultCreditReceiverTest is Test {
    CrossChainFlashVaultCreditReceiver internal recv;
    CreditMockToken internal token;
    address internal router;

    function setUp() public {
        router = makeAddr("ccipRouterCredit");
        recv = new CrossChainFlashVaultCreditReceiver(router);
        token = new CreditMockToken();
    }

    function _msg(address vault, 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("credit-mid"),
            sourceChainSelector: 50,
            sender: abi.encode(address(0x222)),
            data: abi.encode(vault),
            tokenAmounts: amounts
        });
    }

    function test_ccipReceive_creditsVault() public {
        address vaultAddr = address(0xF1A5);
        uint256 amt = 333e18;
        token.mint(address(recv), amt);

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

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

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

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