smom-dbis-138/test/utils/FeeCollectorTest.t.sol

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