// 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"; import {AaveQuotePushFlashReceiver} from "../../contracts/flash/AaveQuotePushFlashReceiver.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; IERC20(address(token)).safeTransfer(to, amount); } } contract MockManagedCycleReceiver is Ownable { using SafeERC20 for IERC20; MockTreasuryToken internal immutable token; uint256 internal simulatedSurplusRaw; constructor(address initialOwner, MockTreasuryToken token_) Ownable(initialOwner) { token = token_; } function setSimulatedSurplusRaw(uint256 amount) external { simulatedSurplusRaw = amount; } function flashQuotePush(address asset, uint256, AaveQuotePushFlashReceiver.QuotePushParams calldata) external { require(asset == address(token), "wrong asset"); if (simulatedSurplusRaw > 0) { token.mint(address(this), simulatedSurplusRaw); } } 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; IERC20(address(token)).safeTransfer(to, amount); } } contract QuotePushTreasuryManagerTest is Test { MockTreasuryToken internal token; MockSweepableReceiver internal receiver; QuotePushTreasuryManager internal manager; MockManagedCycleReceiver internal cycleReceiver; QuotePushTreasuryManager internal cycleManager; 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); cycleReceiver = new MockManagedCycleReceiver(address(this), token); cycleManager = new QuotePushTreasuryManager( address(this), address(cycleReceiver), address(token), OPERATOR, GAS_RECIPIENT, RECYCLE_RECIPIENT, 10_000, 5_000 ); cycleReceiver.transferOwnership(address(cycleManager)); } 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"); } function testRunManagedCycleHarvestsAndDistributesAtomically() public { cycleReceiver.setSimulatedSurplusRaw(80_000); AaveQuotePushFlashReceiver.QuotePushParams memory params = AaveQuotePushFlashReceiver.QuotePushParams({ integration: address(0x1111), pmmPool: address(0x2222), baseToken: address(0x3333), externalUnwinder: address(0x4444), minOutPmm: 1, minOutUnwind: 1, unwindData: "", atomicBridge: AaveQuotePushFlashReceiver.AtomicBridgeParams({ coordinator: address(0), sourceChain: 0, destinationChain: 0, destinationAsset: address(0), bridgeAmount: 0, minDestinationAmount: 0, destinationRecipient: address(0), destinationDeadline: 0, routeId: bytes32(0), settlementMode: bytes32(0), submitCommitment: false }) }); vm.prank(OPERATOR); (uint256 harvested, uint256 gasAmount, uint256 recycleAmount) = cycleManager.runManagedCycle(address(token), 20_000, params, true, 15_000); assertEq(harvested, 70_000, "receiver reserve should stay behind"); assertEq(gasAmount, 15_000, "gas holdback should consume configured slice"); assertEq(recycleAmount, 50_000, "remaining available quote should recycle"); assertEq(token.balanceOf(GAS_RECIPIENT), 15_000, "gas recipient receives holdback"); assertEq(token.balanceOf(RECYCLE_RECIPIENT), 50_000, "recycle recipient receives remaining quote"); assertEq(token.balanceOf(address(cycleReceiver)), 10_000, "receiver reserve is retained"); assertEq(token.balanceOf(address(cycleManager)), 5_000, "manager reserve is retained"); } function testOwnerCanTransferManagedReceiverOwnership() public { address newOwner = address(0xABCD); manager.transferManagedReceiverOwnership(newOwner); assertEq(manager.receiverOwner(), newOwner, "receiver ownership should move to requested address"); assertFalse(manager.isReceiverOwnedByManager(), "manager should no longer own receiver"); } function testRunManagedCycleRevertsWhenManagerDoesNotOwnReceiver() public { AaveQuotePushFlashReceiver.QuotePushParams memory params = AaveQuotePushFlashReceiver.QuotePushParams({ integration: address(0x1111), pmmPool: address(0x2222), baseToken: address(0x3333), externalUnwinder: address(0x4444), minOutPmm: 1, minOutUnwind: 1, unwindData: "", atomicBridge: AaveQuotePushFlashReceiver.AtomicBridgeParams({ coordinator: address(0), sourceChain: 0, destinationChain: 0, destinationAsset: address(0), bridgeAmount: 0, minDestinationAmount: 0, destinationRecipient: address(0), destinationDeadline: 0, routeId: bytes32(0), settlementMode: bytes32(0), submitCommitment: false }) }); cycleManager.transferManagedReceiverOwnership(address(this)); vm.prank(OPERATOR); vm.expectRevert(QuotePushTreasuryManager.ReceiverNotOwnedByManager.selector); cycleManager.runManagedCycle(address(token), 20_000, params, true, 15_000); } }