fix: mocks (cont'd)

src/contracts/components/QWAaveV3.sol

@@ -41,7 +41,7 @@ contract QWAaveV3 is IQWComponent, QWComponentBase {
     function open(uint256 _amount) external override onlyQwManager returns (bool success, uint256 assetAmountReceived) {
         _checkInvestment(_amount);
 
-        token.approve(POOL, _amount);
+        IERC20(INVESTMENT_TOKEN).approve(POOL, _amount);
         IPool(POOL).supply(INVESTMENT_TOKEN, _amount, address(this), 0);
 
         assetAmountReceived = _checkAssetsAny();
@@ -61,7 +61,7 @@ contract QWAaveV3 is IQWComponent, QWComponentBase {
     function close(uint256 _amount) external override onlyQwManager returns (bool success, uint256 tokenAmountReceived) {
         _checkAssets(_amount);
 
-        IPool(POOL).withdraw(INVESTMENT_TOKEN, amountToWithdraw, QW_MANAGER);
+        IPool(POOL).withdraw(INVESTMENT_TOKEN, _amount, QW_MANAGER);
 
         tokenAmountReceived = _checkInvestmentAny();
 

src/contracts/components/QWComponentBase.sol

@@ -54,10 +54,12 @@ abstract contract QWComponentBase is IQWComponent {
     /**
      * @notice Checks whether we've received the proper amount of investment tokens.
      */
-    function _checkInvestment(uint256 _expectedAmount) internal view {
-        if (IERC20(INVESTMENT_TOKEN).balanceOf(address(this)) != _expectedAmount) {
+    function _checkInvestment(uint256 _expectedAmount) internal view returns (uint256) {
+        uint256 balance = IERC20(INVESTMENT_TOKEN).balanceOf(address(this));
+        if (balance != _expectedAmount) {
             revert IncorrectInvestmentTokensReceived(balance);
         }
+        return balance;
     }
 
     /**
@@ -68,7 +70,7 @@ abstract contract QWComponentBase is IQWComponent {
         if (balance == 0) {
             revert IncorrectInvestmentTokensReceived(balance);
         }
-        return balance
+        return balance;
     }
 
     /**

src/contracts/components/QWCompound.sol

@@ -43,7 +43,7 @@ contract QWCompound is IQWComponent, QWComponentBase {
     ) external override onlyQwManager returns (bool success, uint256 assetAmountReceived) {
         _checkInvestment(_amount);
 
-        token.approve(COMET, _amount);
+        IERC20(INVESTMENT_TOKEN).approve(COMET, _amount);
 
         // Perform the supply to Compound.
         IComet(COMET).supplyTo(address(this), INVESTMENT_TOKEN, _amount);
@@ -67,7 +67,7 @@ contract QWCompound is IQWComponent, QWComponentBase {
         _checkAssets(_amount);
 
         // Perform the withdraw from Compound.
-        IComet(COMET).withdrawTo(address(this), INVESTMENT_TOKEN, amountToWithdraw);
+        IComet(COMET).withdrawTo(address(this), INVESTMENT_TOKEN, _amount);
 
         tokenAmountReceived = _checkInvestmentAny();
         success = true;

src/contracts/components/QWUniswapV3Stable.sol

@@ -81,9 +81,10 @@ contract QWUniswapV3Stable is IQWComponent, QWComponentBase, Ownable, IERC721Rec
         (liquidity, amount0, amount1) = increaseLiquidityCurrentRange(_amount);
 
         assetAmountReceived = amount0 + amount1;
-        success = true;
 
         // TODO: Send NFT to QWManager
+
+        success = true;
     }
 
     /**
@@ -101,8 +102,10 @@ contract QWUniswapV3Stable is IQWComponent, QWComponentBase, Ownable, IERC721Rec
 
         (uint256 amount0, uint256 amount1) = decreaseLiquidity();
         tokenAmountReceived = amount0 + amount1;
-        success = true;
+
         // TODO: Transfer NFT back to QWManager.
+
+        success = true;
     }
 
     /**

src/interfaces/IQWRegistry.sol

@@ -14,12 +14,12 @@ interface IQWRegistry {
   function registerComponent(address _component) external;
 
   /**
-   * @notice Checks if a child contract is whitelisted.
-   * @dev Returns true if the specified child contract is whitelisted, otherwise false.
-   * @param _child The address of the child contract to check.
-   * @return A boolean indicating whether the child contract is whitelisted.
+   * @notice Checks if a component contract is whitelisted.
+   * @dev Returns true if the specified component contract is whitelisted, otherwise false.
+   * @param _component The address of the component contract to check.
+   * @return boolean A boolean indicating whether the component contract is whitelisted.
    */
-  function whitelist(address _child) external view returns (bool);
+  function whitelist(address _component) external view returns (bool);
 
   /**
    * @notice Gets the address of the Quant Wealth Manager contract.

test/smock/MockQWManager.sol

@@ -1,48 +1,41 @@
 // SPDX-License-Identifier: APACHE
 pragma solidity ^0.8.0;
 
-import {
-  IERC20, IQWComponent, IQWManager, IQWRegistry, Ownable, QWManager, QWRegistry
-} from '../../src/contracts/QWManager.sol';
+import {IERC20, IQWComponent, IQWManager, IQWRegistry, Ownable, QWManager, QWRegistry} from '../../src/contracts/QWManager.sol';
 import {Test} from 'forge-std/Test.sol';
 
 contract MockQWManager is QWManager, Test {
-  constructor() QWManager() {}
+    constructor() QWManager() {}
 
-  function mock_call_execute(
-    address[] memory _targetQwChild,
-    bytes[] memory _callData,
-    address _tokenAddress,
-    uint256 _amount
-  ) public {
-    vm.mockCall(
-      address(this),
-      abi.encodeWithSignature(
-        'execute(address[],bytes[],address,uint256)', _targetQwChild, _callData, _tokenAddress, _amount
-      ),
-      abi.encode()
-    );
-  }
+    function mock_call_open(OpenBatch[] memory batches) public {
+        vm.mockCall(
+            address(this),
+            abi.encodeWithSignature('open((address,uint256)[])', batches),
+            abi.encode()
+        );
+    }
 
-  function mock_call_close(address[] memory _targetQwChild, bytes[] memory _callData) public {
-    vm.mockCall(
-      address(this), abi.encodeWithSignature('close(address[],bytes[])', _targetQwChild, _callData), abi.encode()
-    );
-  }
+    function mock_call_close(CloseBatch[] memory batches) public {
+        vm.mockCall(
+            address(this),
+            abi.encodeWithSignature('close((address,uint256)[])', batches),
+            abi.encode()
+        );
+    }
 
-  function mock_call_withdraw(address user, address _tokenAddress, uint256 _amount) public {
-    vm.mockCall(
-      address(this),
-      abi.encodeWithSignature('withdraw(address,address,uint256)', user, _tokenAddress, _amount),
-      abi.encode()
-    );
-  }
+    function mock_call_withdraw(address user, address _tokenAddress, uint256 _amount) public {
+        vm.mockCall(
+            address(this),
+            abi.encodeWithSignature('withdraw(address,address,uint256)', user, _tokenAddress, _amount),
+            abi.encode()
+        );
+    }
 
-  function mock_call_receiveFunds(address user, address _tokenAddress, uint256 _amount) public {
-    vm.mockCall(
-      address(this),
-      abi.encodeWithSignature('receiveFunds(address,address,uint256)', user, _tokenAddress, _amount),
-      abi.encode()
-    );
-  }
+    function mock_call_receiveFunds(address user, address _tokenAddress, uint256 _amount) public {
+        vm.mockCall(
+            address(this),
+            abi.encodeWithSignature('receiveFunds(address,address,uint256)', user, _tokenAddress, _amount),
+            abi.encode()
+        );
+    }
 }

test/smock/MockQWRegistry.sol

@@ -5,6 +5,8 @@ import {IQWComponent, IQWRegistry, Ownable, QWRegistry} from '../../src/contract
 import {Test} from 'forge-std/Test.sol';
 
 contract MockQWRegistry is QWRegistry, Test {
+  constructor(address _qwManager) QWRegistry(_qwManager, msg.sender) {}
+
   function set_whitelist(address _key0, bool _value) public {
     whitelist[_key0] = _value;
   }
@@ -13,9 +15,7 @@ contract MockQWRegistry is QWRegistry, Test {
     vm.mockCall(address(this), abi.encodeWithSignature('whitelist(address)', _key0), abi.encode(_value));
   }
 
-  constructor(address _qwManager) QWRegistry(_qwManager, msg.sender) {}
-
-  function mock_call_registerChild(address _child) public {
-    vm.mockCall(address(this), abi.encodeWithSignature('registerChild(address)', _child), abi.encode());
+  function mock_call_registerComponent(address _child) public {
+    vm.mockCall(address(this), abi.encodeWithSignature('registerComponent(address)', _child), abi.encode());
   }
 }

test/smock/components/MockQWAaveV2.sol

@@ -5,17 +5,22 @@ import {IERC20, ILendingPool, IQWComponent, QWAaveV2} from '../../../src/contrac
 import {Test} from 'forge-std/Test.sol';
 
 contract MockQWAaveV2 is QWAaveV2, Test {
-  constructor(address _qwManager, address _lendingPool) QWAaveV2(_qwManager, _lendingPool) {}
+  constructor(
+    address _qwManager,
+    address _investmentToken,
+    address _assetToken,
+    address _pool
+  ) QWAaveV2(_qwManager, _investmentToken, _assetToken, _pool) {}
 
-  function mock_call_create(bytes memory _callData, address _tokenAddress, uint256 _amount, bool success) public {
+  function mock_call_open(uint256 _amount, bool success) public {
     vm.mockCall(
       address(this),
-      abi.encodeWithSignature('create(bytes,address,uint256)', _callData, _tokenAddress, _amount),
+      abi.encodeWithSignature('open(uint256)', _amount),
       abi.encode(success)
     );
   }
 
-  function mock_call_close(bytes memory _callData, bool success) public {
-    vm.mockCall(address(this), abi.encodeWithSignature('close(bytes)', _callData), abi.encode(success));
+  function mock_call_close(uint256 _amount, bool success) public {
+    vm.mockCall(address(this), abi.encodeWithSignature('close(uint256)', _amount), abi.encode(success));
   }
 }

test/smock/components/MockQWAaveV3.sol

@@ -5,17 +5,22 @@ import {IERC20, IPool, IQWComponent, QWAaveV3} from '../../../src/contracts/comp
 import {Test} from 'forge-std/Test.sol';
 
 contract MockQWAaveV3 is QWAaveV3, Test {
-  constructor(address _qwManager, address _pool) QWAaveV3(_qwManager, _pool) {}
+  constructor(
+    address _qwManager,
+    address _investmentToken,
+    address _assetToken,
+    address _pool
+  ) QWAaveV3(_qwManager, _investmentToken, _assetToken, _pool) {}
 
-  function mock_call_create(bytes memory _callData, address _tokenAddress, uint256 _amount, bool success) public {
+  function mock_call_open(uint256 _amount, bool success) public {
     vm.mockCall(
       address(this),
-      abi.encodeWithSignature('create(bytes,address,uint256)', _callData, _tokenAddress, _amount),
+      abi.encodeWithSignature('open(uint256)', _amount),
       abi.encode(success)
     );
   }
 
-  function mock_call_close(bytes memory _callData, bool success) public {
-    vm.mockCall(address(this), abi.encodeWithSignature('close(bytes)', _callData), abi.encode(success));
+  function mock_call_close(uint256 _amount, bool success) public {
+    vm.mockCall(address(this), abi.encodeWithSignature('close(uint256)', _amount), abi.encode(success));
   }
 }

test/smock/components/MockQWCompound.sol

@@ -12,15 +12,15 @@ contract MockQWCompound is QWCompound, Test {
     address _comet
   ) QWCompound(_qwManager, _investmentToken, _assetToken, _comet) {}
 
-  function mock_call_create(bytes memory _callData, address _tokenAddress, uint256 _amount, bool success) public {
+  function mock_call_open(uint256 _amount, bool success) public {
     vm.mockCall(
       address(this),
-      abi.encodeWithSignature('create(bytes,address,uint256)', _callData, _tokenAddress, _amount),
+      abi.encodeWithSignature('open(uint256)', _amount),
       abi.encode(success)
     );
   }
 
-  function mock_call_close(bytes memory _callData, bool success) public {
-    vm.mockCall(address(this), abi.encodeWithSignature('close(bytes)', _callData), abi.encode(success));
+  function mock_call_close(uint256 _amount, bool success) public {
+    vm.mockCall(address(this), abi.encodeWithSignature('close(uint256)', _amount), abi.encode(success));
   }
 }