feat: add working withdraw

src/contracts/Grateful.sol

@@ -8,10 +8,19 @@ import {IGrateful} from "interfaces/IGrateful.sol";
 import {AaveV3ERC4626, IPool, IRewardsController} from "yield-daddy/aave-v3/AaveV3ERC4626.sol";
 
 contract Grateful is IGrateful, Ownable2Step {
+  // @inheritdoc IGrateful
   IPool public aavePool;
+
+  // @inheritdoc IGrateful
   mapping(address => bool) public tokensWhitelisted;
+
+  // @inheritdoc IGrateful
   mapping(address => bool) public yieldingFunds;
+
+  // @inheritdoc IGrateful
   mapping(address => AaveV3ERC4626) public vaults;
+
+  // @inheritdoc IGrateful
   mapping(address => mapping(address => uint256)) public shares;
 
   modifier onlyWhenTokenWhitelisted(address _token) {
@@ -29,9 +38,16 @@ contract Grateful is IGrateful, Ownable2Step {
     }
   }
 
+  /// @inheritdoc IGrateful
+  function addToken(address _token) external onlyOwner {
+    tokensWhitelisted[_token] = true;
+    IERC20(_token).approve(address(aavePool), type(uint256).max);
+  }
+
   // @inheritdoc IGrateful
   function addVault(address _token, address _vault) external onlyWhenTokenWhitelisted(_token) onlyOwner {
     vaults[_token] = AaveV3ERC4626(_vault);
+    IERC20(_token).approve(address(_vault), type(uint256).max);
   }
 
   // @inheritdoc IGrateful
@@ -51,6 +67,17 @@ contract Grateful is IGrateful, Ownable2Step {
     }
   }
 
+  // @inheritdoc IGrateful
+  function withdraw(address _token) external onlyWhenTokenWhitelisted(_token) {
+    AaveV3ERC4626 vault = vaults[_token];
+    if (address(vault) == address(0)) {
+      revert Grateful_VaultNotSet();
+    }
+    uint256 _shares = shares[msg.sender][_token];
+    shares[msg.sender][_token] = 0;
+    vault.redeem(_shares, msg.sender, address(this));
+  }
+
   // @inheritdoc IGrateful
   function switchYieldingFunds() external {
     yieldingFunds[msg.sender] = !yieldingFunds[msg.sender];

src/interfaces/IGrateful.sol

@@ -75,6 +75,12 @@ interface IGrateful {
                             LOGIC
   //////////////////////////////////////////////////////////////*/
 
+  /**
+   * @notice Adds a token to the whitelist
+   * @param _token Address of the token to be added to the whitelist
+   */
+  function addToken(address _token) external;
+
   /**
    * @notice Makes a payment to a merchant
    * @param _merchant Address of the merchant receiving payment
@@ -83,6 +89,12 @@ interface IGrateful {
    */
   function pay(address _merchant, address _token, uint256 _amount) external;
 
+  /**
+   * @notice Withdraws funds from the vault
+   * @param _token Address of the token being withdrawn
+   */
+  function withdraw(address _token) external;
+
   /**
    * @notice Switch the preference of the merchant to yield funds or not
    */

test/integration/Grateful.t.sol

@@ -1,17 +1,16 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.26;
 
-import {console} from "forge-std/console.sol";
 import {IntegrationBase} from "test/integration/IntegrationBase.sol";
 
 contract IntegrationGreeter is IntegrationBase {
   function test_Payment() public {
     vm.startPrank(_daiWhale);
-    _dai.approve(address(_grateful), 100);
-    _grateful.pay(_merchant, address(_dai), 100);
+    _dai.approve(address(_grateful), _amount);
+    _grateful.pay(_merchant, address(_dai), _amount);
     vm.stopPrank();
 
-    assertEq(_dai.balanceOf(_merchant), 100);
+    assertEq(_dai.balanceOf(_merchant), _amount);
   }
 
   function test_PaymentYieldingFunds() public {
@@ -23,8 +22,15 @@ contract IntegrationGreeter is IntegrationBase {
     assertEq(_grateful.yieldingFunds(_merchant), true);
 
     vm.startPrank(_daiWhale);
-    _dai.approve(address(_grateful), 100);
-    _grateful.pay(_merchant, address(_dai), 100);
+    _dai.approve(address(_grateful), _amount);
+    _grateful.pay(_merchant, address(_dai), _amount);
     vm.stopPrank();
+
+    vm.warp(block.timestamp + 60 days);
+
+    vm.prank(_merchant);
+    _grateful.withdraw(address(_dai));
+
+    assertGt(_dai.balanceOf(_merchant), _amount);
   }
 }

test/integration/IntegrationBase.sol

@@ -5,6 +5,8 @@ import {Grateful, IGrateful} from "contracts/Grateful.sol";
 import {Test} from "forge-std/Test.sol";
 
 import {IERC20} from "forge-std/interfaces/IERC20.sol";
+
+import {ERC4626} from "solmate/mixins/ERC4626.sol";
 import {ERC20} from "solmate/tokens/ERC20.sol";
 import {AaveV3ERC4626, IPool, IRewardsController} from "yield-daddy/aave-v3/AaveV3ERC4626.sol";
 
@@ -19,21 +21,24 @@ contract IntegrationBase is Test {
   address[] internal _tokens;
   IERC20 internal _dai = IERC20(0x6B175474E89094C44Da98b954EedeAC495271d0F);
   address _aDai = 0x018008bfb33d285247A21d44E50697654f754e63;
+  address _rewardsController = 0x8164Cc65827dcFe994AB23944CBC90e0aa80bFcb;
   IPool internal _aavePool = IPool(0x87870Bca3F3fD6335C3F4ce8392D69350B4fA4E2);
   IGrateful internal _grateful;
-  AaveV3ERC4626 internal _vault = new AaveV3ERC4626(
-    ERC20(address(_dai)),
-    ERC20(_aDai),
-    _aavePool,
-    address(0),
-    IRewardsController(address(_dai)) // TODO: put real rewards controller
-  );
+  AaveV3ERC4626 internal _vault;
+  uint256 internal _amount = 1e25;
 
   function setUp() public {
+    vm.startPrank(_owner);
     vm.createSelectFork(vm.rpcUrl("mainnet"), _FORK_BLOCK);
-    vm.prank(_owner);
+    _vault = new AaveV3ERC4626(
+      ERC20(address(_dai)), ERC20(_aDai), _aavePool, address(0), IRewardsController(_rewardsController)
+    );
+    vm.label(address(_vault), "Vault");
     _tokens = new address[](1);
     _tokens[0] = address(_dai);
     _grateful = new Grateful(_tokens, _aavePool);
+    vm.label(address(_grateful), "Grateful");
+    _grateful.addVault(address(_dai), address(_vault));
+    vm.stopPrank();
   }
 }