From 8aa1a04e37b999d95ce6b6199703bd30398923e8 Mon Sep 17 00:00:00 2001
From: 0xchin <alanracciatti1220@gmail.com>
Date: Thu, 14 Nov 2024 12:15:02 -0300
Subject: [PATCH] test: add assets calculation

---
 src/contracts/Grateful.sol      | 17 +++++++++--------
 src/interfaces/IGrateful.sol    |  5 +++--
 test/integration/Grateful.t.sol |  8 +++++++-
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/contracts/Grateful.sol b/src/contracts/Grateful.sol
index a237401..e673040 100644
--- a/src/contracts/Grateful.sol
+++ b/src/contracts/Grateful.sol
@@ -8,11 +8,12 @@ import {IERC20} from "@openzeppelin/contracts/interfaces/IERC20.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
 import {OneTime} from "contracts/OneTime.sol";
+import {AaveV3Vault} from "contracts/vaults/AaveV3Vault.sol";
 
 import {IGrateful} from "interfaces/IGrateful.sol";
 
 import {Bytes32AddressLib} from "solmate/utils/Bytes32AddressLib.sol";
-import {AaveV3ERC4626, IPool} from "yield-daddy/aave-v3/AaveV3ERC4626.sol";
+import {IPool} from "yield-daddy/aave-v3/AaveV3ERC4626.sol";
 
 /**
  * @title Grateful Contract
@@ -33,7 +34,7 @@ contract Grateful is IGrateful, Ownable2Step {
   mapping(address => bool) public tokensWhitelisted;
 
   /// @inheritdoc IGrateful
-  mapping(address => AaveV3ERC4626) public vaults;
+  mapping(address => AaveV3Vault) public vaults;
 
   /// @inheritdoc IGrateful
   mapping(address => mapping(address => uint256)) public shares;
@@ -133,7 +134,7 @@ contract Grateful is IGrateful, Ownable2Step {
 
   /// @inheritdoc IGrateful
   function addVault(address _token, address _vault) external onlyOwner onlyWhenTokenWhitelisted(_token) {
-    vaults[_token] = AaveV3ERC4626(_vault);
+    vaults[_token] = AaveV3Vault(_vault);
     IERC20(_token).safeIncreaseAllowance(address(_vault), type(uint256).max);
   }
 
@@ -200,7 +201,7 @@ contract Grateful is IGrateful, Ownable2Step {
   function withdraw(
     address _token
   ) external onlyWhenTokenWhitelisted(_token) {
-    AaveV3ERC4626 vault = vaults[_token];
+    AaveV3Vault vault = vaults[_token];
     if (address(vault) == address(0)) {
       revert Grateful_usdcVaultNotSet();
     }
@@ -211,7 +212,7 @@ contract Grateful is IGrateful, Ownable2Step {
 
   /// @inheritdoc IGrateful
   function withdraw(address _token, uint256 _assets) external onlyWhenTokenWhitelisted(_token) {
-    AaveV3ERC4626 vault = vaults[_token];
+    AaveV3Vault vault = vaults[_token];
     if (address(vault) == address(0)) {
       revert Grateful_usdcVaultNotSet();
     }
@@ -231,7 +232,7 @@ contract Grateful is IGrateful, Ownable2Step {
     uint256 tokensLength = _tokens.length;
     for (uint256 i = 0; i < tokensLength; i++) {
       address _token = _tokens[i];
-      AaveV3ERC4626 vault = vaults[_token];
+      AaveV3Vault vault = vaults[_token];
       if (address(vault) == address(0)) {
         revert Grateful_usdcVaultNotSet();
       }
@@ -255,7 +256,7 @@ contract Grateful is IGrateful, Ownable2Step {
     for (uint256 i = 0; i < tokensLength; i++) {
       address _token = _tokens[i];
       uint256 _assetsToWithdraw = _assets[i];
-      AaveV3ERC4626 vault = vaults[_token];
+      AaveV3Vault vault = vaults[_token];
       if (address(vault) == address(0)) {
         revert Grateful_usdcVaultNotSet();
       }
@@ -338,7 +339,7 @@ contract Grateful is IGrateful, Ownable2Step {
     IERC20(_token).safeTransfer(owner(), _amount - amountWithFee);
 
     if (_yieldFunds) {
-      AaveV3ERC4626 vault = vaults[_token];
+      AaveV3Vault vault = vaults[_token];
       if (address(vault) == address(0)) {
         IERC20(_token).safeTransfer(_merchant, amountWithFee);
       } else {
diff --git a/src/interfaces/IGrateful.sol b/src/interfaces/IGrateful.sol
index 914a471..a2a7539 100644
--- a/src/interfaces/IGrateful.sol
+++ b/src/interfaces/IGrateful.sol
@@ -4,7 +4,8 @@ pragma solidity 0.8.26;
 import {OneTime} from "contracts/OneTime.sol";
 
 import {IERC20} from "@openzeppelin/contracts/interfaces/IERC20.sol";
-import {AaveV3ERC4626, IPool} from "yield-daddy/aave-v3/AaveV3ERC4626.sol";
+import {AaveV3Vault} from "contracts/vaults/AaveV3Vault.sol";
+import {IPool} from "yield-daddy/aave-v3/AaveV3ERC4626.sol";
 
 /**
  * @title Grateful Contract Interface
@@ -108,7 +109,7 @@ interface IGrateful {
   /// @return Address of the vault contract.
   function vaults(
     address _token
-  ) external view returns (AaveV3ERC4626);
+  ) external view returns (AaveV3Vault);
 
   /// @notice Returns the amount of shares for a merchant.
   /// @param _merchant Address of the merchant.
diff --git a/test/integration/Grateful.t.sol b/test/integration/Grateful.t.sol
index 0fccc96..c00369f 100644
--- a/test/integration/Grateful.t.sol
+++ b/test/integration/Grateful.t.sol
@@ -2,7 +2,7 @@
 pragma solidity 0.8.26;
 
 import {OneTime} from "contracts/OneTime.sol";
-import {IntegrationBase} from "test/integration/IntegrationBase.sol";
+import {AaveV3Vault, IntegrationBase} from "test/integration/IntegrationBase.sol";
 
 contract IntegrationGreeter is IntegrationBase {
   function _approveAndPay(address payer, address merchant, uint256 amount, bool _yieldFunds) internal {
@@ -24,9 +24,15 @@ contract IntegrationGreeter is IntegrationBase {
 
     vm.warp(block.timestamp + 1 days);
 
+    // Get total assets
+    AaveV3Vault _vault = _grateful.vaults(address(_usdc));
+    uint256 _shares = _grateful.shares(_merchant, address(_usdc));
+    uint256 _assets = _vault.convertToAssets(_shares);
+
     vm.prank(_merchant);
     _grateful.withdraw(address(_usdc));
 
+    assertEq(_assets, _usdc.balanceOf(_merchant));
     assertGt(_usdc.balanceOf(_merchant), _grateful.applyFee(_merchant, _AMOUNT_USDC));
   }