make wrapModuleV2 and setValuer mutable

contracts/adapters/TargetWeightWrapExtension.sol

@@ -29,6 +29,7 @@ import { BaseExtension } from "../lib/BaseExtension.sol";
 import { Position } from "../lib/Position.sol";
 import { PreciseUnitMath } from "../lib/PreciseUnitMath.sol";
 import { IBaseManager } from "../interfaces/IBaseManager.sol";
+import { IController } from "../interfaces/IController.sol";
 import { ISetToken } from "../interfaces/ISetToken.sol";
 import { ISetValuer } from "../interfaces/ISetValuer.sol";
 import { IWrapModuleV2 } from "../interfaces/IWrapModuleV2.sol";
@@ -76,15 +77,18 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
         TargetWeightWrapParams[] executionParams
     );
     event RebalancePaused();
+    event WrapModuleUpdated(address indexed wrapModule);
+    event SetValuerUpdated(address indexed setValuer);
 
     /* ========== Immutables ========= */
 
     ISetToken public immutable setToken;
-    IWrapModuleV2 public immutable wrapModule;
-    ISetValuer public immutable setValuer;
 
     /* ========== State Variables ========= */
 
+    IWrapModuleV2 public wrapModule;
+    ISetValuer public setValuer;
+
     bool public isRebalancingActive;
     bool public isRebalanceOpen;
 
@@ -114,10 +118,14 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
         bool _isRebalanceOpen
     ) public BaseExtension(_manager) {
         manager = _manager;
-        setToken = manager.setToken();
         wrapModule = _wrapModule;
         setValuer = _setValuer;
         isRebalanceOpen = _isRebalanceOpen;
+
+        ISetToken setToken_ = manager.setToken();
+        setToken = setToken_;
+        _setWrapModule(setToken_, _wrapModule);
+        _setSetValuer(setToken_, _setValuer);
     }
 
     /* ========== Rebalance Functions ========== */
@@ -254,6 +262,24 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
         emit RebalanceAccessUpdated(_isRebalanceOpen);
     }
 
+    /**
+     * @notice Sets the WrapModule contract used for wrapping and unwrapping assets.
+     * @dev This function can only be called by the operator.
+     * @param _wrapModule Address of the WrapModuleV2 contract.
+     */
+    function setWrapModule(IWrapModuleV2 _wrapModule) external onlyOperator {
+        _setWrapModule(setToken, _wrapModule);
+    }
+
+    /**
+     * @notice Sets the SetValuer contract used for calculating valuations and weights.
+     * @dev This function can only be called by the operator.
+     * @param _setValuer Address of the SetValuer contract.
+     */
+    function setSetValuer(ISetValuer _setValuer) external onlyOperator {
+        _setSetValuer(setToken, _setValuer);
+    }
+
     /**
      * @notice Initializes the Set Token within the Wrap Module.
      * @dev This function can only be called by the operator.
@@ -370,6 +396,30 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
         return rebalanceInfo.targetAssets;
     }
 
+    /* ========== Internal Functions ========== */
+
+    /**
+     * Sets the WrapModuleV2 contract used for wrapping and unwrapping assets.
+     * @param _setToken Address of the SetToken contract.
+     * @param _wrapModule Address of the WrapModuleV2 contract.
+     */
+    function _setWrapModule(ISetToken _setToken, IWrapModuleV2 _wrapModule) internal {
+        require(_setToken.moduleStates(address(_wrapModule)) == ISetToken.ModuleState.PENDING, "WrapModuleV2 not pending");
+        wrapModule = _wrapModule;
+        emit WrapModuleUpdated(address(_wrapModule));
+    }
+
+    /**
+     * Sets the SetValuer contract used for calculating valuations and weights.
+     * @param _setToken Address of the SetToken contract.
+     * @param _setValuer Address of the SetValuer contract.
+     */
+    function _setSetValuer(ISetToken _setToken, ISetValuer _setValuer) internal {
+        require(IController(_setToken.controller()).isResource(address(_setValuer)), "SetValuer not approved by controller");
+        setValuer = _setValuer;
+        emit SetValuerUpdated(address(_setValuer));
+    }
+
     /* ============== Modifier Helpers ===============
      * Internal functions used to reduce bytecode size
      */

contracts/interfaces/IController.sol

@@ -20,6 +20,7 @@ interface IController {
     function isModule(address _module) external view returns(bool);
     function isSet(address _setToken) external view returns(bool);
     function isSystemContract(address _contractAddress) external view returns (bool);
+    function isResource(address _resource) external view returns(bool);
     function resourceId(uint256 _id) external view returns(address);
     function owner() external view returns(address);
     function addFactory(address _factory) external;

contracts/interfaces/ISetToken.sol

@@ -72,6 +72,8 @@ interface ISetToken is IERC20 {
 
     /* ============ Functions ============ */
 
+    function controller() external view returns (address);
+
     function addComponent(address _component) external;
     function removeComponent(address _component) external;
     function editDefaultPositionUnit(address _component, int256 _realUnit) external;
@@ -115,4 +117,4 @@ interface ISetToken is IERC20 {
     function isInitializedModule(address _module) external view returns(bool);
     function isPendingModule(address _module) external view returns(bool);
     function isLocked() external view returns (bool);
-}
+}

test/adapters/targetWeightWrapExtension.spec.ts

@@ -2,7 +2,7 @@ import "module-alias/register";
 
 import DeployHelper from "@utils/deploys";
 import { SetFixture } from "@utils/fixtures";
-import { SetToken } from "@utils/contracts/setV2";
+import { SetToken, WrapModuleV2 } from "@utils/contracts/setV2";
 import {
   BaseManagerV2,
   TargetWeightWrapExtension,
@@ -17,6 +17,7 @@ import {
   getRandomAccount,
   getWaffleExpect,
   preciseDiv,
+  getRandomAddress,
 } from "@utils/index";
 import { ADDRESS_ZERO, MAX_UINT_256, ZERO, ZERO_BYTES } from "@utils/constants";
 import { BigNumber, ContractTransaction } from "ethers";
@@ -361,6 +362,92 @@ describe("TargetWeightWrapExtension", async () => {
         });
       });
 
+      describe("#setIsWrapModule", () => {
+        let subjectCaller: Account;
+        let subjectWrapModule: WrapModuleV2;
+
+        beforeEach(async () => {
+          subjectWrapModule = await deployer.setV2.deployWrapModuleV2(
+            setV2Setup.controller.address,
+            setV2Setup.weth.address,
+          );
+          await setV2Setup.controller.addModule(subjectWrapModule.address);
+          await baseManager.connect(operator.wallet).addModule(subjectWrapModule.address);
+          subjectCaller = operator;
+        });
+
+        function subject() {
+          return targetWeightWrapExtension.connect(subjectCaller.wallet).setWrapModule(subjectWrapModule.address);
+        }
+
+        it("should set the WrapModuleV2 correctly", async () => {
+          await subject();
+          expect(await targetWeightWrapExtension.wrapModule()).to.eq(subjectWrapModule.address);
+        });
+
+        context("when the module is not pending", async () => {
+          beforeEach(async () => {
+            subjectWrapModule = await deployer.setV2.deployWrapModuleV2(
+              setV2Setup.controller.address,
+              setV2Setup.weth.address,
+            );
+          });
+
+          it("should revert", async () => {
+            await expect(subject()).to.be.revertedWith("WrapModuleV2 not pending");
+          });
+        });
+
+        context("when the caller is not the operator", async () => {
+          beforeEach(async () => {
+            subjectCaller = feeRecipient;
+          });
+
+          it("should revert", async () => {
+            await expect(subject()).to.be.revertedWith("Must be operator");
+          });
+        });
+      });
+
+      describe("#setSetValuer", () => {
+        let subjectCaller: Account;
+        let subjectSetValuer: Address;
+
+        beforeEach(async () => {
+          subjectCaller = operator;
+          subjectSetValuer = setV2Setup.integrationRegistry.address;
+        });
+
+        function subject() {
+          return targetWeightWrapExtension.connect(subjectCaller.wallet).setSetValuer(subjectSetValuer);
+        }
+
+        it("should set the SetValuer correctly", async () => {
+          await subject();
+          expect(await targetWeightWrapExtension.setValuer()).to.eq(subjectSetValuer);
+        });
+
+        context("when the setValuer is not approved on the controller", async () => {
+          beforeEach(async () => {
+            subjectSetValuer = await getRandomAddress();
+          });
+
+          it("should revert", async () => {
+            await expect(subject()).to.be.revertedWith("SetValuer not approved by controller");
+          });
+        });
+
+        context("when the caller is not the operator", async () => {
+          beforeEach(async () => {
+            subjectCaller = feeRecipient;
+          });
+
+          it("should revert", async () => {
+            await expect(subject()).to.be.revertedWith("Must be operator");
+          });
+        });
+      });
+
       context("when targets are set", async () => {
         beforeEach(async () => {
           await targetWeightWrapExtension.connect(operator.wallet).setTargetWeights(