// SPDX-License-Identifier: MIT pragma solidity ^0.8.19; import {Test} from "forge-std/Test.sol"; import {FeeCollector} from "../../contracts/utils/FeeCollector.sol"; import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; contract MockERC20 is ERC20 { constructor() ERC20("Mock Token", "MOCK") { _mint(msg.sender, 1000000 * 10**18); } } contract FeeCollectorTest is Test { FeeCollector public collector; MockERC20 public token; address public admin; address public feeManager; address public recipient1; address public recipient2; function setUp() public { admin = address(this); feeManager = address(this); recipient1 = address(0x1); recipient2 = address(0x2); collector = new FeeCollector(admin); token = new MockERC20(); } function testCollectETHFees() public { uint256 amount = 1 ether; (bool success, ) = address(collector).call{value: amount}( abi.encodeWithSignature("collectFees(address,uint256)", address(0), amount) ); assertTrue(success); assertEq(collector.getBalance(address(0)), amount); assertEq(collector.getTotalCollected(address(0)), amount); } function testCollectTokenFees() public { uint256 amount = 1000 * 10**18; token.approve(address(collector), amount); collector.collectFees(address(token), amount); assertEq(collector.getBalance(address(token)), amount); assertEq(collector.getTotalCollected(address(token)), amount); } function testAddFeeRecipient() public { collector.addFeeRecipient(address(0), recipient1, 5000); // 50% collector.addFeeRecipient(address(0), recipient2, 5000); // 50% FeeCollector.FeeRecipient[] memory recipients = collector.getFeeRecipients(address(0)); assertEq(recipients.length, 2); assertEq(recipients[0].recipient, recipient1); assertEq(recipients[0].shareBps, 5000); } function testDistributeETHFees() public { uint256 amount = 1 ether; // Collect fees (bool success, ) = address(collector).call{value: amount}( abi.encodeWithSignature("collectFees(address,uint256)", address(0), amount) ); assertTrue(success); // Add recipients collector.addFeeRecipient(address(0), recipient1, 5000); // 50% collector.addFeeRecipient(address(0), recipient2, 5000); // 50% // Distribute collector.distributeFees(address(0)); assertEq(recipient1.balance, 0.5 ether); assertEq(recipient2.balance, 0.5 ether); assertEq(collector.getBalance(address(0)), 0); } function testDistributeTokenFees() public { uint256 amount = 1000 * 10**18; // Collect fees token.approve(address(collector), amount); collector.collectFees(address(token), amount); // Add recipients collector.addFeeRecipient(address(token), recipient1, 6000); // 60% collector.addFeeRecipient(address(token), recipient2, 4000); // 40% // Distribute collector.distributeFees(address(token)); assertEq(token.balanceOf(recipient1), 600 * 10**18); assertEq(token.balanceOf(recipient2), 400 * 10**18); assertEq(collector.getBalance(address(token)), 0); } function testRemoveFeeRecipient() public { collector.addFeeRecipient(address(0), recipient1, 5000); collector.addFeeRecipient(address(0), recipient2, 5000); collector.removeFeeRecipient(address(0), recipient1); FeeCollector.FeeRecipient[] memory recipients = collector.getFeeRecipients(address(0)); assertEq(recipients.length, 1); assertEq(recipients[0].recipient, recipient2); } function testEmergencyWithdraw() public { uint256 amount = 1 ether; // Collect fees (bool success, ) = address(collector).call{value: amount}( abi.encodeWithSignature("collectFees(address,uint256)", address(0), amount) ); assertTrue(success); // Emergency withdraw collector.emergencyWithdraw(address(0), recipient1, amount); assertEq(recipient1.balance, amount); assertEq(collector.getBalance(address(0)), 0); } }