smom-dbis-138/test/flash/QuotePushTreasuryManager.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 {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import {QuotePushTreasuryManager} from "../../contracts/flash/QuotePushTreasuryManager.sol";

contract MockTreasuryToken is ERC20 {
    constructor() ERC20("Mock Treasury Token", "MTT") {}

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

contract MockSweepableReceiver is Ownable {
    using SafeERC20 for IERC20;

    IERC20 internal immutable token;

    constructor(address initialOwner, IERC20 token_) Ownable(initialOwner) {
        token = token_;
    }

    function quoteSurplusBalance(address quoteToken, uint256 reserveRetained) external view returns (uint256 surplus) {
        require(quoteToken == address(token), "wrong token");
        uint256 balance = token.balanceOf(address(this));
        if (balance > reserveRetained) {
            surplus = balance - reserveRetained;
        }
    }

    function sweepQuoteSurplus(address quoteToken, address to, uint256 reserveRetained)
        external
        onlyOwner
        returns (uint256 amount)
    {
        require(quoteToken == address(token), "wrong token");
        uint256 balance = token.balanceOf(address(this));
        require(balance > reserveRetained, "nothing to sweep");
        amount = balance - reserveRetained;
        token.safeTransfer(to, amount);
    }
}

contract QuotePushTreasuryManagerTest is Test {
    MockTreasuryToken internal token;
    MockSweepableReceiver internal receiver;
    QuotePushTreasuryManager internal manager;

    address internal constant OPERATOR = address(0xBEEF);
    address internal constant GAS_RECIPIENT = address(0xCAFE);
    address internal constant RECYCLE_RECIPIENT = address(0xF00D);

    function setUp() public {
        token = new MockTreasuryToken();
        receiver = new MockSweepableReceiver(address(this), token);
        manager = new QuotePushTreasuryManager(
            address(this), address(receiver), address(token), OPERATOR, GAS_RECIPIENT, RECYCLE_RECIPIENT, 200_000, 100_000
        );

        receiver.transferOwnership(address(manager));
        token.mint(address(receiver), 1_000_000);
    }

    function testHarvestReceiverSurplusKeepsReceiverReserve() public {
        assertTrue(manager.isReceiverOwnedByManager(), "manager should own receiver");
        assertEq(manager.receiverSweepableQuote(), 800_000, "sweepable quote should exclude receiver reserve");

        uint256 harvested = manager.harvestReceiverSurplus();

        assertEq(harvested, 800_000, "harvested amount should match receiver surplus");
        assertEq(token.balanceOf(address(receiver)), 200_000, "receiver should keep configured reserve");
        assertEq(manager.quoteBalance(), 800_000, "manager should receive harvested quote");
        assertEq(manager.availableQuote(), 700_000, "manager reserve should be retained");
    }

    function testOperatorCanDistributeConfiguredRecipients() public {
        manager.harvestReceiverSurplus();

        vm.prank(OPERATOR);
        manager.distributeToConfiguredRecipients(250_000, 300_000);

        assertEq(token.balanceOf(GAS_RECIPIENT), 250_000, "gas recipient should receive configured amount");
        assertEq(token.balanceOf(RECYCLE_RECIPIENT), 300_000, "recycle recipient should receive configured amount");
        assertEq(manager.quoteBalance(), 250_000, "manager should retain the undistributed quote");
        assertEq(manager.availableQuote(), 150_000, "manager reserve should still be protected");
    }

    function testDistributeRevertsWhenRequestedAmountExceedsAvailable() public {
        manager.harvestReceiverSurplus();

        vm.prank(OPERATOR);
        vm.expectRevert();
        manager.distributeToConfiguredRecipients(500_000, 250_001);
    }

    function testOwnerCanRescueNonQuoteToken() public {
        MockTreasuryToken other = new MockTreasuryToken();
        other.mint(address(manager), 42);

        manager.rescueToken(address(other), address(this), 42);

        assertEq(other.balanceOf(address(this)), 42, "owner should be able to rescue unrelated tokens");
    }
}