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