feat: Private function to allow both access control and ownable two s…

…teps
LooksRare · Jul 4, 2024 · 6f67949 · 6f67949
1 parent 9ca51b7
commit 6f67949
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 17 deletions.
diff --git a/contracts/BlastYield.sol b/contracts/BlastYield.sol
@@ -2,7 +2,6 @@
 
 pragma solidity ^0.8.20;
 
-import {OwnableTwoSteps} from "./OwnableTwoSteps.sol";
 import {IBlast, YieldMode as IBlast__YieldMode, GasMode as IBlast__GasMode} from "./interfaces/IBlast.sol";
 import {IBlastPoints} from "./interfaces/IBlastPoints.sol";
 import {IERC20Rebasing, YieldMode as IERC20Rebasing__YieldMode} from "./interfaces/IERC20Rebasing.sol";
@@ -12,7 +11,7 @@ import {IERC20Rebasing, YieldMode as IERC20Rebasing__YieldMode} from "./interfac
  * @notice This contract is a base contract for future contracts that wish to claim Blast WETH or USDB yield to inherit from.
  * @author LooksRare protocol team (👀,💎)
  */
-contract BlastYield is OwnableTwoSteps {
+contract BlastYield {
     address public immutable WETH;
     address public immutable USDB;
 
@@ -31,7 +30,7 @@ contract BlastYield is OwnableTwoSteps {
         address _owner,
         address _usdb,
         address _weth
-    ) OwnableTwoSteps(_owner) {
+    ) {
         WETH = _weth;
         USDB = _usdb;
 
@@ -46,7 +45,7 @@ contract BlastYield is OwnableTwoSteps {
      * @param wethReceiver The receiver of WETH.
      * @param usdbReceiver The receiver of USDB.
      */
-    function claim(address wethReceiver, address usdbReceiver) external virtual onlyOwner {
+    function _claim(address wethReceiver, address usdbReceiver) internal virtual {
         uint256 claimableWETH = IERC20Rebasing(WETH).getClaimableAmount(address(this));
         if (claimableWETH != 0) {
             IERC20Rebasing(WETH).claim(wethReceiver, claimableWETH);

diff --git a/test/foundry/BlastYield.t.sol b/test/foundry/BlastYield.t.sol
@@ -3,6 +3,7 @@
 pragma solidity ^0.8.20;
 
 import {IOwnableTwoSteps} from "../../contracts/interfaces/IOwnableTwoSteps.sol";
+import {OwnableTwoSteps} from "../../contracts/OwnableTwoSteps.sol";
 import {BlastYield} from "../../contracts/BlastYield.sol";
 import {Test} from "../../lib/forge-std/src/Test.sol";
 import {YieldMode as IBlast__YieldMode, GasMode as IBlast__GasMode} from "../../contracts/interfaces/IBlast.sol";
@@ -13,12 +14,27 @@ import {MockPoints} from "../mock/MockBlastPoints.sol";
 import {MockWETH} from "../mock/MockBlastWETH.sol";
 import {MockYield} from "../mock/MockBlastYield.sol";
 
-contract BlastYield_Test is Test {
+contract BlastYieldOwnableTwoSteps is BlastYield, OwnableTwoSteps {
+    constructor(
+        address _blast,
+        address _blastPoints,
+        address _blastPointsOperator,
+        address _owner,
+        address _usdb,
+        address _weth
+    ) BlastYield(_blast, _blastPoints, _blastPointsOperator, _owner, _usdb, _weth) OwnableTwoSteps(_owner) {}
+
+    function claim(address wethReceiver, address usdbReceiver) public onlyOwner {
+        _claim(wethReceiver, usdbReceiver);
+    }
+}
+
+contract BlastYieldOwnableTwoSteps_Test is Test {
     MockWETH private weth;
     MockERC20 private usdb;
     MockYield private mockYield;
     MockPoints private mockPoints;
-    BlastYield private blastYield;
+    BlastYieldOwnableTwoSteps private blastYieldOwnableTwoSteps;
 
     address public owner = address(69);
     address public operator = address(420);
@@ -32,40 +48,46 @@ contract BlastYield_Test is Test {
         weth = new MockWETH();
         usdb = new MockERC20("USDB", "USDB");
         mockYield = new MockYield();
-        blastYield = new BlastYield(address(mockYield), BLAST_POINTS, operator, owner, address(usdb), address(weth));
+        blastYieldOwnableTwoSteps = new BlastYieldOwnableTwoSteps(
+            address(mockYield),
+            BLAST_POINTS,
+            operator,
+            owner,
+            address(usdb),
+            address(weth)
+        );
     }
 
     function test_setUpState() public {
-        assertEq(blastYield.WETH(), address(weth));
-        assertEq(blastYield.USDB(), address(usdb));
-        assertEq(blastYield.owner(), owner);
+        assertEq(blastYieldOwnableTwoSteps.WETH(), address(weth));
+        assertEq(blastYieldOwnableTwoSteps.USDB(), address(usdb));
 
         (IBlast__YieldMode yieldMode, IBlast__GasMode gasMode, address governor) = mockYield.config(
-            address(blastYield)
+            address(blastYieldOwnableTwoSteps)
         );
         assertEq(uint8(yieldMode), uint8(IBlast__YieldMode.CLAIMABLE));
         assertEq(uint8(gasMode), uint8(IBlast__GasMode.CLAIMABLE));
         assertEq(governor, owner);
 
-        IERC20Rebasing__YieldMode wethYieldMode = weth.yieldMode(address(blastYield));
+        IERC20Rebasing__YieldMode wethYieldMode = weth.yieldMode(address(blastYieldOwnableTwoSteps));
         assertEq(uint8(wethYieldMode), uint8(IERC20Rebasing__YieldMode.CLAIMABLE));
 
-        IERC20Rebasing__YieldMode usdbYieldMode = usdb.yieldMode(address(blastYield));
+        IERC20Rebasing__YieldMode usdbYieldMode = usdb.yieldMode(address(blastYieldOwnableTwoSteps));
         assertEq(uint8(usdbYieldMode), uint8(IERC20Rebasing__YieldMode.CLAIMABLE));
     }
 
     function test_claim() public asPrankedUser(owner) {
-        blastYield.claim(TREASURY, TREASURY);
+        blastYieldOwnableTwoSteps.claim(TREASURY, TREASURY);
 
-        assertEq(weth.balanceOf(address(blastYield)), 0);
-        assertEq(usdb.balanceOf(address(blastYield)), 0);
+        assertEq(weth.balanceOf(address(blastYieldOwnableTwoSteps)), 0);
+        assertEq(usdb.balanceOf(address(blastYieldOwnableTwoSteps)), 0);
         assertEq(weth.balanceOf(TREASURY), 1 ether);
         assertEq(usdb.balanceOf(TREASURY), 1 ether);
     }
 
     function test_claim_RevertIf_NotOwner() public asPrankedUser(user1) {
         vm.expectRevert(IOwnableTwoSteps.NotOwner.selector);
-        blastYield.claim(TREASURY, TREASURY);
+        blastYieldOwnableTwoSteps.claim(TREASURY, TREASURY);
     }
 
     modifier asPrankedUser(address user) {