refactor: adjusted close fn with lpAsset

QuantWealth · May 18, 2024 · 91f6714 · 91f6714
1 parent 1595786
commit 91f6714
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/src/contracts/QWManager.sol b/src/contracts/QWManager.sol
@@ -52,8 +52,7 @@ contract QWManager is IQWManager {
       }
 
       IERC20 token = IERC20(_tokenAddress);
-      token.approve(_targetQwChild[i], _amount);
-      token.transferFrom(address(this), address(_targetQwChild[i]), _amount);
+      token.approve(address(_targetQwChild[i]), _amount);
 
       (bool success) = IQWChild(_targetQwChild[i]).create(_callData[i], _tokenAddress, _amount);
       if (!success) {
@@ -74,6 +73,13 @@ contract QWManager is IQWManager {
     }
 
     for (uint256 i = 0; i < _targetQwChild.length; i++) {
+      (, address lpAsset, uint256 amount) = abi.decode(
+          _callData[i],
+          (address, address, uint256)
+      );
+      IERC20 token = IERC20(lpAsset);
+      token.approve(address(_targetQwChild[i]), amount);
+
       (bool success) = IQWChild(_targetQwChild[i]).close(_callData[i]);
       if (!success) {
         revert CallFailed();

diff --git a/src/contracts/child/QWAave.sol b/src/contracts/child/QWAave.sol
@@ -46,6 +46,7 @@ contract QWAave is IQWChild {
         }
 
         IERC20 token = IERC20(_tokenAddress);
+        token.transferFrom(QW_MANAGER, address(this), _amount);
         token.approve(POOL, _amount);
 
         IPool(POOL).supply(_tokenAddress, _amount, QW_MANAGER, 0);
@@ -65,11 +66,15 @@ contract QWAave is IQWChild {
             revert InvalidCallData();
         }
 
-        (address asset, uint256 amount) = abi.decode(
+        (address asset, address lpAsset, uint256 amount) = abi.decode(
             _callData,
-            (address, uint256)
+            (address, address, uint256)
         );
 
+        IERC20 token = IERC20(lpAsset);
+        token.transferFrom(QW_MANAGER, address(this), amount);
+        token.approve(POOL, amount);
+
         IPool(POOL).withdraw(asset, amount, QW_MANAGER);
         return true;
     }

diff --git a/test/integration/child/QWAave.t.sol b/test/integration/child/QWAave.t.sol
@@ -51,8 +51,11 @@ contract AaveIntegration is IntegrationBase {
         test_Create_Aave();
 
         vm.startPrank(_daiWhale);
-        uint256 amount = 1_000_000 ether;
-        bytes memory callData = abi.encode(_dai, amount);
+        uint256 amount = _aDai.balanceOf(address(_qwManager));
+        bytes memory callData = abi.encode(address(_dai), address(_aDai), amount);
+
+        uint256 aDaiBalanceBefore = _aDai.balanceOf(address(_qwManager));
+        uint256 daiBalanceBefore = _dai.balanceOf(address(_qwManager));
 
         // Create dynamic arrays with one element each
         address[] memory targetQWChild = new address[](1);
@@ -63,6 +66,13 @@ contract AaveIntegration is IntegrationBase {
 
         // bug: qwAave contract need to have aDai tokens and approve them to aave pool contract to withdraw
         _qwManager.close(targetQWChild, callDataArr);
+
+        uint256 aDaiBalanceAfter = _aDai.balanceOf(address(_qwManager));
+        uint256 daiBalanceAfter = _dai.balanceOf(address(_qwManager));
+
+        assertGe(daiBalanceAfter - daiBalanceBefore, amount);
+        assertEq(aDaiBalanceBefore - aDaiBalanceAfter, amount);
+        assertEq(aDaiBalanceAfter, 0);
         vm.stopPrank();
     }
 }