smom-dbis-138/test/flash/QuotePushFlashWorkflowBorrower.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 {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {IERC3156FlashBorrower} from "@openzeppelin/contracts/interfaces/IERC3156FlashBorrower.sol";
import {SimpleERC3156FlashVault} from "../../contracts/flash/SimpleERC3156FlashVault.sol";
import {
    QuotePushFlashWorkflowBorrower,
    IDODOQuotePushSwapExactIn,
    IExternalUnwinder
} from "../../contracts/flash/QuotePushFlashWorkflowBorrower.sol";

contract MockQuoteToken is ERC20 {
    constructor(string memory n, string memory s) ERC20(n, s) {}

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

/// @notice Simple PMM-like mock: quote in, base out at a fixed ratio.
contract MockQuotePushIntegration is IDODOQuotePushSwapExactIn {
    IERC20 public immutable quote;
    IERC20 public immutable base;
    uint256 public immutable numerator;
    uint256 public immutable denominator;

    constructor(IERC20 quote_, IERC20 base_, uint256 numerator_, uint256 denominator_) {
        quote = quote_;
        base = base_;
        numerator = numerator_;
        denominator = denominator_;
    }

    function swapExactIn(address, address tokenIn, uint256 amountIn, uint256 minAmountOut)
        external
        override
        returns (uint256 amountOut)
    {
        require(tokenIn == address(quote), "quote only");
        IERC20(tokenIn).transferFrom(msg.sender, address(this), amountIn);
        amountOut = amountIn * numerator / denominator;
        require(amountOut >= minAmountOut, "min pmm");
        IERC20(address(base)).transfer(msg.sender, amountOut);
    }
}

/// @notice Simple external unwinder: base in, quote out at a fixed ratio.
contract MockExternalUnwinder is IExternalUnwinder {
    IERC20 public immutable quote;
    IERC20 public immutable base;
    uint256 public immutable numerator;
    uint256 public immutable denominator;

    constructor(IERC20 base_, IERC20 quote_, uint256 numerator_, uint256 denominator_) {
        base = base_;
        quote = quote_;
        numerator = numerator_;
        denominator = denominator_;
    }

    function unwind(address tokenIn, address tokenOut, uint256 amountIn, uint256 minAmountOut, bytes calldata)
        external
        override
        returns (uint256 amountOut)
    {
        require(tokenIn == address(base), "base only");
        require(tokenOut == address(quote), "quote only");
        IERC20(tokenIn).transferFrom(msg.sender, address(this), amountIn);
        amountOut = amountIn * numerator / denominator;
        require(amountOut >= minAmountOut, "min unwind");
        IERC20(address(quote)).transfer(msg.sender, amountOut);
    }
}

contract QuotePushFlashWorkflowBorrowerTest is Test {
    SimpleERC3156FlashVault internal vault;
    MockQuoteToken internal usdc;
    MockQuoteToken internal cwusdc;
    MockQuotePushIntegration internal pmm;
    MockExternalUnwinder internal unwinder;
    QuotePushFlashWorkflowBorrower internal borrower;

    address internal owner = address(0xA11);
    address internal user = address(0xB22);

    function setUp() public {
        vm.startPrank(owner);
        vault = new SimpleERC3156FlashVault(owner, 5);
        usdc = new MockQuoteToken("USDC", "USDC");
        cwusdc = new MockQuoteToken("cWUSDC", "cWUSDC");
        usdc.mint(address(vault), 1_000_000e6);
        vault.setTokenSupported(address(usdc), true);
        vm.stopPrank();

        // PMM leg: slightly under 1:1, like a fee-bearing quote->base purchase.
        pmm = new MockQuotePushIntegration(usdc, cwusdc, 9997, 10000);
        // External unwind: profitable 1.12x unwind.
        unwinder = new MockExternalUnwinder(cwusdc, usdc, 112, 100);
        borrower = new QuotePushFlashWorkflowBorrower(address(vault));

        cwusdc.mint(address(pmm), 10_000_000e6);
        usdc.mint(address(unwinder), 10_000_000e6);
    }

    function test_quotePushRoundTrip_repayAndRetainSurplus() public {
        uint256 amount = 4_145_894; // 4.145894 USDC
        uint256 fee = vault.flashFee(address(usdc), amount);

        QuotePushFlashWorkflowBorrower.QuotePushParams memory p = QuotePushFlashWorkflowBorrower.QuotePushParams({
            integration: address(pmm),
            pool: address(0x69776fc607e9edA8042e320e7e43f54d06c68f0E),
            baseToken: address(cwusdc),
            externalUnwinder: address(unwinder),
            minOutPmm: 4_100_000,
            minOutUnwind: amount + fee,
            unwindData: bytes("")
        });

        uint256 feesBefore = vault.totalFeesCollected(address(usdc));
        vm.prank(user);
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), amount, abi.encode(p));

        uint256 feesAfter = vault.totalFeesCollected(address(usdc));
        uint256 borrowerSurplus = usdc.balanceOf(address(borrower));

        assertEq(feesAfter - feesBefore, fee, "flash fee collected");
        assertGt(borrowerSurplus, 0, "surplus retained");
        assertEq(cwusdc.balanceOf(address(borrower)), 0, "base fully unwound");
    }

    function test_quotePushRoundTrip_revert_whenUnwindCannotRepay() public {
        uint256 amount = 4_145_894; // 4.145894 USDC

        // Replace profitable unwinder with a losing one.
        MockExternalUnwinder badUnwinder = new MockExternalUnwinder(cwusdc, usdc, 95, 100);
        usdc.mint(address(badUnwinder), 10_000_000e6);

        QuotePushFlashWorkflowBorrower.QuotePushParams memory p = QuotePushFlashWorkflowBorrower.QuotePushParams({
            integration: address(pmm),
            pool: address(0x69776fc607e9edA8042e320e7e43f54d06c68f0E),
            baseToken: address(cwusdc),
            externalUnwinder: address(badUnwinder),
            minOutPmm: 4_100_000,
            minOutUnwind: 1,
            unwindData: bytes("")
        });

        vm.expectRevert(QuotePushFlashWorkflowBorrower.InsufficientToRepay.selector);
        vm.prank(user);
        vault.flashLoan(IERC3156FlashBorrower(address(borrower)), address(usdc), amount, abi.encode(p));
    }
}