Files
smom-dbis-138/test/flash/QuotePushTreasuryManager.t.sol
2026-04-13 21:37:33 -07:00

240 lines
9.8 KiB
Solidity

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