// 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)); } }