smom-dbis-138/test/flash/AaveQuotePushFlashReceiver.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 {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");
    }
}