// 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 {AaveQuotePushFlashReceiver} from "../../contracts/flash/AaveQuotePushFlashReceiver.sol"; contract MockSweepToken is ERC20 { constructor() ERC20("Mock Sweep Token", "MST") {} function mint(address to, uint256 amount) external { _mint(to, amount); } } contract AaveQuotePushFlashReceiverTest is Test { AaveQuotePushFlashReceiver internal receiver; MockSweepToken internal token; function setUp() public { receiver = new AaveQuotePushFlashReceiver(address(0x1234), address(this)); token = new MockSweepToken(); token.mint(address(receiver), 1_000_000); } function testOwnerCanSweepQuoteSurplusAndKeepReserve() public { uint256 reserveRetained = 250_000; uint256 sweepable = receiver.quoteSurplusBalance(address(token), reserveRetained); assertEq(sweepable, 750_000, "sweepable amount should exclude retained reserve"); receiver.sweepQuoteSurplus(address(token), address(this), reserveRetained); assertEq(token.balanceOf(address(receiver)), reserveRetained, "receiver should keep reserve"); assertEq(token.balanceOf(address(this)), 750_000, "owner should receive swept surplus"); } function testNonOwnerCannotSweep() public { vm.prank(address(0xBEEF)); vm.expectRevert(); receiver.sweepQuoteSurplus(address(token), address(0xBEEF), 0); } function testOwnerCanSweepExplicitTokenAmount() public { receiver.sweepToken(address(token), address(0xCAFE), 123_456); assertEq(token.balanceOf(address(receiver)), 876_544, "receiver balance should decrease"); assertEq(token.balanceOf(address(0xCAFE)), 123_456, "recipient should receive exact amount"); } }