Test cases for FeeCollector

contracts/core/TTUFeeCollector.sol

@@ -10,48 +10,63 @@ contract TTUFeeCollector is ITTUFeeCollector, Ownable {
     using SafeERC20 for IERC20;
 
     uint256 public constant BIPS_PRECISION = 10 ** 4;
-    uint256 public constant MAX_FEE = 10 ** 3;
+    uint256 public constant MAX_FEE_BIPS = 10 ** 3;
 
     uint256 public defaultFeesBips;
     mapping(address => uint256) internal _customFeesBips;
+    mapping(address => uint256) internal _customFeesFixed;
 
     constructor() Ownable(_msgSender()) {}
 
     function withdrawFee(IERC20 token, uint256 amount) external onlyOwner {
         token.safeTransfer(owner(), amount);
     }
 
-    function setDefaultFee(uint256 bips) external onlyOwner {
-        if (bips > MAX_FEE) revert FeesTooHigh();
+    function setDefaultFeeBips(uint256 bips) external onlyOwner {
+        if (bips > MAX_FEE_BIPS) revert FeesTooHigh();
         defaultFeesBips = bips;
-        emit DefaultFeeSet(bips);
+        emit DefaultFeeSetBips(bips);
     }
 
-    // @dev Setting bips to BIPS_PRECISION means 0 fees
-    function setCustomFee(
+    // @dev Setting bips to MAX_FEE_BIPS means 0 fees, so technically the MAX_FEE_BIPS is MAX_FEE_BIPS - 1
+    function setCustomFeeBips(
         address unlockerAddress,
         uint256 bips
     ) external onlyOwner {
-        if (bips > MAX_FEE) revert FeesTooHigh();
+        if (bips > MAX_FEE_BIPS) revert FeesTooHigh();
         _customFeesBips[unlockerAddress] = bips;
-        emit CustomFeeSet(unlockerAddress, bips);
+        emit CustomFeeSetBips(unlockerAddress, bips);
+    }
+
+    function setCustomFeeFixed(
+        address unlockerAddress,
+        uint256 fixedFee
+    ) external onlyOwner {
+        // Not capping a MAX_FEE for fixed fee since it cannot apply equally to all token values
+        _customFeesFixed[unlockerAddress] = fixedFee;
+        emit CustomFeeSetFixed(unlockerAddress, fixedFee);
     }
 
     function getFee(
         address unlockerAddress,
         uint256 tokenTransferred
     ) external view override returns (uint256 tokensCollected) {
+        uint256 feeFixed = _customFeesFixed[unlockerAddress];
+        // If there is a fixed fee, return that fee immediately without checking bips fee
+        if (feeFixed > 0) {
+            return feeFixed;
+        }
         uint256 feeBips = _customFeesBips[unlockerAddress];
         if (feeBips == 0) {
             feeBips = defaultFeesBips;
-        } else if (feeBips == BIPS_PRECISION) {
+        } else if (feeBips == MAX_FEE_BIPS) {
             feeBips = 0;
         }
         tokensCollected = (tokenTransferred * feeBips) / BIPS_PRECISION;
     }
 
     function version() external pure returns (string memory) {
-        return "2.1.0";
+        return "2.6.0";
     }
 
     function transferOwnership(

contracts/interfaces/ITTUFeeCollector.sol

@@ -10,16 +10,17 @@ import {IVersionable} from "./IVersionable.sol";
  * @dev This contract handles TokenTable service fee calculation.
  */
 interface ITTUFeeCollector is IOwnable, IVersionable {
-    event DefaultFeeSet(uint256 bips);
-    event CustomFeeSet(address unlockerAddress, uint256 bips);
+    event DefaultFeeSetBips(uint256 bips);
+    event CustomFeeSetBips(address unlockerAddress, uint256 bips);
+    event CustomFeeSetFixed(address unlockerAddress, uint256 fixedFee);
 
     /**
      * @dev 0xc9034e18
      */
     error FeesTooHigh();
 
     /**
-     * @notice Returns the amount of fees to collect.
+     * @notice Returns the amount of fees to collect. A fixed fee will always override a dynamic fee.
      * @param unlockerAddress The address of the Unlocker. Used to fetch pricing.
      * @param tokenTransferred The number of tokens transferred.
      * @return tokensCollected The number of tokens to collect as fees.

test/TTUFeeCollector.t.sol

@@ -0,0 +1,168 @@
+// SPDX-License-Identifier: UNLICENSED
+// solhint-disable ordering
+pragma solidity ^0.8.20;
+
+import {Test, console} from "forge-std/Test.sol";
+import {TTUFeeCollector} from "../contracts/core/TTUFeeCollector.sol";
+import {MockERC20} from "../contracts/mock/MockERC20.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+
+contract TTUFeeCollectorTest is Test {
+    using SafeERC20 for IERC20;
+
+    TTUFeeCollector public instance;
+    IERC20 public mockErc20;
+
+    event DefaultFeeSetBips(uint256 bips);
+    event CustomFeeSetBips(address unlockerAddress, uint256 bips);
+    event CustomFeeSetFixed(address unlockerAddress, uint256 fixedFee);
+
+    error FeesTooHigh();
+
+    // Ownable
+    error OwnableUnauthorizedAccount(address account);
+
+    function setUp() public {
+        instance = new TTUFeeCollector();
+        mockErc20 = IERC20(address(new MockERC20()));
+    }
+
+    function testFuzz_setDefaultFeeBips_success(uint256 bips) public {
+        vm.assume(bips <= instance.MAX_FEE_BIPS());
+        vm.expectEmit();
+        emit DefaultFeeSetBips(bips);
+        instance.setDefaultFeeBips(bips);
+    }
+
+    function testFuzz_setDefaultFeeBips_fail_notOwner(
+        address notOwner,
+        uint256 bips
+    ) public {
+        vm.assume(bips <= instance.MAX_FEE_BIPS() && notOwner != address(this));
+        vm.prank(notOwner);
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                OwnableUnauthorizedAccount.selector,
+                notOwner
+            )
+        );
+        instance.setDefaultFeeBips(bips);
+    }
+
+    function testFuzz_setDefaultFeeBips_fail_feesTooHigh(uint256 bips) public {
+        vm.assume(bips > instance.MAX_FEE_BIPS());
+        vm.expectRevert(abi.encodeWithSelector(FeesTooHigh.selector));
+        instance.setDefaultFeeBips(bips);
+    }
+
+    function testFuzz_setCustomFeeBips_success(
+        address unlockerAddress,
+        uint256 bips
+    ) public {
+        vm.assume(bips <= instance.MAX_FEE_BIPS());
+        vm.expectEmit();
+        emit CustomFeeSetBips(unlockerAddress, bips);
+        instance.setCustomFeeBips(unlockerAddress, bips);
+    }
+
+    function testFuzz_setCustomFeeBips_fail_notOwner(
+        address notOwner,
+        address unlockerAddress,
+        uint256 bips
+    ) public {
+        vm.assume(bips <= instance.MAX_FEE_BIPS() && notOwner != address(this));
+        vm.prank(notOwner);
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                OwnableUnauthorizedAccount.selector,
+                notOwner
+            )
+        );
+        instance.setCustomFeeBips(unlockerAddress, bips);
+    }
+
+    function testFuzz_setCustomFeeBips_fail_feesTooHigh(
+        address unlockerAddress,
+        uint256 bips
+    ) public {
+        vm.assume(bips > instance.MAX_FEE_BIPS());
+        vm.expectRevert(abi.encodeWithSelector(FeesTooHigh.selector));
+        instance.setCustomFeeBips(unlockerAddress, bips);
+    }
+
+    function testFuzz_getFee(
+        uint16 defaultFeeBips,
+        address customFeeUnlockerAddress,
+        uint256 customFeeFixed,
+        uint16 customFeeBips,
+        uint128 tokenTransferred
+    ) public {
+        vm.assume(
+            defaultFeeBips <= instance.MAX_FEE_BIPS() &&
+                customFeeBips <= instance.MAX_FEE_BIPS()
+        );
+        instance.setDefaultFeeBips(defaultFeeBips);
+        instance.setCustomFeeBips(customFeeUnlockerAddress, customFeeBips);
+        instance.setCustomFeeFixed(customFeeUnlockerAddress, customFeeFixed);
+        uint256 fees = instance.getFee(
+            customFeeUnlockerAddress,
+            tokenTransferred
+        );
+        if (customFeeFixed > 0) {
+            assertEq(fees, customFeeFixed);
+        } else if (customFeeBips > 0) {
+            if (customFeeBips == instance.MAX_FEE_BIPS()) {
+                assertEq(fees, 0);
+            } else {
+                assertEq(
+                    fees,
+                    (uint256(tokenTransferred) * customFeeBips) /
+                        instance.BIPS_PRECISION()
+                );
+            }
+        } else {
+            assertEq(
+                fees,
+                (uint256(tokenTransferred) * defaultFeeBips) /
+                    instance.BIPS_PRECISION()
+            );
+        }
+    }
+
+    function testFuzz_withdrawFee_success(uint256 amount) public {
+        MockERC20(address(mockErc20)).mint(address(this), amount);
+        mockErc20.safeTransfer(address(instance), amount);
+        uint256 balanceBefore = mockErc20.balanceOf(address(this));
+        instance.withdrawFee(IERC20(mockErc20), amount);
+        assertEq(mockErc20.balanceOf(address(this)) - balanceBefore, amount);
+    }
+
+    function testFuzz_withdrawFee_fail_notOwner(
+        address notOwner,
+        uint256 amount
+    ) public {
+        vm.assume(notOwner != address(this));
+        MockERC20(address(mockErc20)).mint(address(this), amount);
+        mockErc20.safeTransfer(address(instance), amount);
+        vm.prank(notOwner);
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                OwnableUnauthorizedAccount.selector,
+                notOwner
+            )
+        );
+        instance.withdrawFee(IERC20(mockErc20), amount);
+    }
+
+    function testFuzz_withdrawFee_fail_wrongToken(
+        address notMockErc20,
+        uint256 amount
+    ) public {
+        vm.assume(notMockErc20 != address(mockErc20));
+        MockERC20(address(mockErc20)).mint(address(this), amount);
+        mockErc20.safeTransfer(address(instance), amount);
+        vm.expectRevert();
+        instance.withdrawFee(IERC20(notMockErc20), amount);
+    }
+}