add tests for claimYield

Author: acollette

src/treasury/TreasuryV1.sol

@@ -84,14 +84,8 @@ contract TreasuryV1 is StorageV1 {
     function claimYield(uint256 amount, address receiver) external onlyOwner {
         if (amount > availableYield) revert YieldTooLow();
 
-        if (amount == type(uint256).max) {
-            uint256 availableYield_ = availableYield;
-            availableYield = 0;
-            IERC20(EURE).safeTransfer(receiver, availableYield_);
-        } else {
-            availableYield -= amount;
-            IERC20(EURE).safeTransfer(receiver, amount);
-        }
+        availableYield -= amount;
+        IERC20(EURE).safeTransfer(receiver, amount);
     }
 
     /**

test/fuzz/Treasury/ClaimYield.fuzz.t.sol

@@ -0,0 +1,78 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.22;
+
+import {Treasury_Fuzz_Test} from "./_Treasury.fuzz.t.sol";
+
+import {FixedPointMathLib} from "../../../lib/solmate/src/utils/FixedPointMathLib.sol";
+import {IERC20} from "../../../lib/openzeppelin-contracts/contracts/interfaces/IERC20.sol";
+import {ILocker} from "../../../src/lockers/interfaces/ILocker.sol";
+import {TreasuryV1} from "../../../src/treasury/TreasuryV1.sol";
+
+/**
+ * @notice Fuzz tests for the function "claimYield" of contract "Treasury".
+ */
+contract ClaimYield_Treasury_Fuzz_Test is Treasury_Fuzz_Test {
+    using FixedPointMathLib for uint256;
+    /* ///////////////////////////////////////////////////////////////
+                              SETUP
+    /////////////////////////////////////////////////////////////// */
+
+    function setUp() public override {
+        Treasury_Fuzz_Test.setUp();
+    }
+
+    /*//////////////////////////////////////////////////////////////
+                              TESTS
+    //////////////////////////////////////////////////////////////*/
+    function testFuzz_Revert_claimYield_NotOwner(address random, uint256 amount, address receiver) public {
+        vm.assume(random != users.dao);
+
+        vm.startPrank(random);
+        bytes memory expectedError = abi.encodeWithSelector(TreasuryV1.OnlyOwner.selector);
+        vm.expectRevert(expectedError);
+        treasury.claimYield(amount, receiver);
+        vm.stopPrank();
+    }
+
+    function testFuzz_Revert_claimYield_YieldTooLow(uint256 totalYield, uint256 amountToClaim, address receiver)
+        public
+    {
+        // Given: availableYield is lower than yield to claim.
+        vm.assume(amountToClaim > 1);
+        totalYield = bound(totalYield, 0, amountToClaim - 1);
+
+        // And: availableYield is set.
+        treasury.setAvailableYield(totalYield);
+
+        // When: Claiming yield it should revert.
+        vm.startPrank(users.dao);
+        vm.expectRevert(TreasuryV1.YieldTooLow.selector);
+        treasury.claimYield(amountToClaim, receiver);
+        vm.stopPrank();
+    }
+
+    function testFuzz_Success_claimYield(uint256 totalYield, uint256 amountToClaim, address receiver) public {
+        // Given: availableYield is lower than yield to claim.
+        // And: Amount to claim should not be equal to max type(uint256).max.
+        vm.assume(amountToClaim > 0);
+        vm.assume(amountToClaim < type(uint256).max - 2);
+        totalYield = bound(totalYield, amountToClaim + 1, type(uint256).max - 1);
+
+        // And: availableYield is set.
+        treasury.setAvailableYield(totalYield);
+
+        // And: Yield is available in the Treasury
+        EURE.mint(address(treasury), totalYield);
+
+        // When: Claiming yield
+        vm.prank(users.dao);
+        treasury.claimYield(amountToClaim, receiver);
+
+        // Then: Yield should have been sent to receiver
+        // And: Available yield should have been lowered.
+        assertEq(EURE.balanceOf(receiver), amountToClaim);
+        assertEq(EURE.balanceOf(address(treasury)), totalYield - amountToClaim);
+        assertEq(treasury.availableYield(), totalYield - amountToClaim);
+        assert(treasury.availableYield() > 0);
+    }
+}

test/fuzz/Treasury/Initialize.fuzz.t.sol

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.22;
+
+import {Treasury_Fuzz_Test} from "./_Treasury.fuzz.t.sol";
+
+import {FixedPointMathLib} from "../../../lib/solmate/src/utils/FixedPointMathLib.sol";
+import {IERC20} from "../../../lib/openzeppelin-contracts/contracts/interfaces/IERC20.sol";
+import {ILocker} from "../../../src/lockers/interfaces/ILocker.sol";
+import {TreasuryV1} from "../../../src/treasury/TreasuryV1.sol";
+import {TreasuryV1Extension} from "../../utils/extensions/TreasuryV1Extension.sol";
+
+/**
+ * @notice Fuzz tests for the function "initialize" of contract "Treasury".
+ */
+contract Initialize_Treasury_Fuzz_Test is Treasury_Fuzz_Test {
+    using FixedPointMathLib for uint256;
+    /* ///////////////////////////////////////////////////////////////
+                              SETUP
+    /////////////////////////////////////////////////////////////// */
+
+    function setUp() public override {
+        Treasury_Fuzz_Test.setUp();
+    }
+
+    /*//////////////////////////////////////////////////////////////
+                              TESTS
+    //////////////////////////////////////////////////////////////*/
+    function testFuzz_Revert_initialize_AlreadyInitialized() public {
+        vm.expectRevert(TreasuryV1.AlreadyInitialized.selector);
+        treasury.initialize(address(EURE));
+    }
+
+    function testFuzz_Success_initialize(address random) public {
+        TreasuryV1Extension treasury_;
+        vm.startPrank(random);
+        treasury_ = new TreasuryV1Extension();
+        treasury_.initialize(address(EURE));
+
+        assertEq(treasury_.owner(), random);
+        assertEq(treasury_.EURE(), address(EURE));
+    }
+}

test/utils/extensions/TreasuryV1Extension.sol

@@ -20,4 +20,8 @@ contract TreasuryV1Extension is TreasuryV1 {
     function getWeightsLength() public view returns (uint256 length) {
         length = lockersWeights.length;
     }
+
+    function setAvailableYield(uint256 yield) public {
+        availableYield = yield;
+    }
 }