feat: update previewWithdraw to return the deposit ID results (#2)

superical · Jul 12, 2023 · 11e9e99 · 11e9e99
1 parent 4779f9c
commit 11e9e99
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -26,7 +26,7 @@ represent the amount of their deposited asset.
 The `SimpleTimeLockVault` contract is a ready-to-use vault contract, providing the standard
 time-lock vault functionality such as `deposit`, `withdraw` and `batchWithdraw`.
 
-## Custom Vaults
+## Custom Time-Lock Vaults
 
 Alternatively, you can create your own Solidity vault contract by extending the `TimeLockVault`
 abstract contract and implementing your desired logic.

diff --git a/contracts/TimeLockVault.sol b/contracts/TimeLockVault.sol
@@ -165,22 +165,25 @@ abstract contract TimeLockVault is ERC20Upgradeable, ITimeLockVault, TimeLockVau
 
     /**
      * @notice Preview the amount of asset that can be redeemed from the payment IDs
+     * @dev Returns the total amount and an array of boolean indicating whether the specified deposit ID is included
+     * in the calculation
      */
     function previewWithdraw(
         address recipient,
         uint256[] calldata depositIds
-    ) external view virtual returns (uint256) {
-        uint256 totalAmount = 0;
+    ) external view virtual returns (uint256 totalAmount, bool[] memory depositIdsIncluded) {
+        totalAmount = 0;
+        depositIdsIncluded = new bool[](depositIds.length);
         for (uint256 i = 0; i < depositIds.length; i++) {
             DepositInfo memory deposit_ = getDeposit(depositIds[i]);
             if (
                 isDepositActive(recipient, depositIds[i]) &&
                 deposit_.redeemTimestamp <= block.timestamp
             ) {
                 totalAmount += deposit_.amount;
+                depositIdsIncluded[i] = true;
             }
         }
-        return totalAmount;
     }
 
     function _deposit(

diff --git a/contracts/interfaces/ITimeLockVault.sol b/contracts/interfaces/ITimeLockVault.sol
@@ -25,7 +25,7 @@ interface ITimeLockVault is IERC20MetadataUpgradeable {
     function previewWithdraw(
         address recipient,
         uint256[] calldata depositIds
-    ) external view returns (uint256);
+    ) external view returns (uint256 totalAmount, bool[] memory depositIdsIncluded);
 
     function totalActiveDepositsOf(address recipient) external view returns (uint256);
 

diff --git a/test/SimpleTimeLockVault/simple-time-lock-vault-withdraw.spec.ts b/test/SimpleTimeLockVault/simple-time-lock-vault-withdraw.spec.ts
@@ -184,6 +184,11 @@ describe("Simple Time-Lock Vault", () => {
       let recipient1: SignerWithAddress;
       let recipient2: SignerWithAddress;
 
+      let recipient1Amount: BigNumber;
+      let depositIds1Included: boolean[];
+      let recipient2Amount: BigNumber;
+      let depositIds2Included: boolean[];
+
       const oneDayInterval = 3600 * 24;
 
       beforeEach(async () => {
@@ -213,11 +218,11 @@ describe("Simple Time-Lock Vault", () => {
       describe("Withdrawal Preview", () => {
         describe("When all deposit IDs are not matured yet", () => {
           it("should return zero as the amount", async () => {
-            const recipient1Amount = await simpleVaultContract.previewWithdraw(
+            const [recipient1Amount] = await simpleVaultContract.previewWithdraw(
               recipient1.address,
               [0, 2, 4],
             );
-            const recipient2Amount = await simpleVaultContract.previewWithdraw(
+            const [recipient2Amount] = await simpleVaultContract.previewWithdraw(
               recipient2.address,
               [1, 3],
             );
@@ -228,55 +233,75 @@ describe("Simple Time-Lock Vault", () => {
         });
 
         describe("When some of the deposit IDs are matured", () => {
-          it("should return the total amount of the specified deposit IDs that are matured for withdrawal", async () => {
+          beforeEach(async () => {
             // Deposit ID 4 for recipient 1 should not be matured yet
             await time.increase(oneDayInterval * 6);
 
-            const recipient1Amount = await simpleVaultContract.previewWithdraw(
+            [recipient1Amount, depositIds1Included] = await simpleVaultContract.previewWithdraw(
               recipient1.address,
               [0, 2, 4],
             );
-            const recipient2Amount = await simpleVaultContract.previewWithdraw(
+            [recipient2Amount, depositIds2Included] = await simpleVaultContract.previewWithdraw(
               recipient2.address,
               [1, 3],
             );
+          });
 
+          it("should return the total amount of the specified deposit IDs that are matured for withdrawal", async () => {
             expect(recipient1Amount).to.equal(parseEther("200"));
             expect(recipient2Amount).to.equal(parseEther("200"));
           });
+
+          it("should return the included result for the specified deposit IDs that are matured for withdrawal", async () => {
+            expect(depositIds1Included).to.deep.equal([true, true, false]);
+            expect(depositIds2Included).to.deep.equal([true, true]);
+          });
         });
 
         describe("When some of the deposit IDs are inactive deposits", () => {
-          it("should return only the total amount of the deposit IDs that are active", async () => {
+          beforeEach(async () => {
             await time.increase(oneDayInterval * 6);
             await simpleVaultContract.connect(recipient2).withdraw(1);
 
-            const recipient2Amount = await simpleVaultContract.previewWithdraw(
+            [recipient2Amount, depositIds2Included] = await simpleVaultContract.previewWithdraw(
               recipient2.address,
               [1, 3],
             );
+          });
 
+          it("should return only the total amount of the deposit IDs that are active", async () => {
             expect(recipient2Amount).to.equal(parseEther("100"));
           });
+
+          it("should return the included result for the specified deposit IDs that are active", async () => {
+            expect(depositIds2Included).to.deep.equal([false, true]);
+          });
         });
 
         describe("When all the deposit IDs are matured", () => {
-          it("should return the total amount of the specified deposit IDs", async () => {
+          beforeEach(async () => {
             // 7 days later... All payments matured.
             await time.increase(oneDayInterval * 7);
 
-            const recipient1Amount = await simpleVaultContract.previewWithdraw(
+            [recipient1Amount, depositIds1Included] = await simpleVaultContract.previewWithdraw(
               recipient1.address,
               [0, 2, 4],
             );
-            const recipient2Amount = await simpleVaultContract.previewWithdraw(
+            [recipient2Amount, depositIds2Included] = await simpleVaultContract.previewWithdraw(
               recipient2.address,
               [1, 3],
             );
+          });
 
+          it("should return the total amount of the specified deposit IDs", async () => {
             expect(recipient1Amount).to.equal(parseEther("300"));
             expect(recipient2Amount).to.equal(parseEther("200"));
           });
+
+          it("should return the included result for the specified deposit IDs of all matured deposits", async () => {
+            expect(depositIds1Included).to.deep.equal([true, true, true]);
+            expect(depositIds2Included).to.deep.equal([true, true]);
+          });
         });
 
         it("should revert if withdraws a premature deposit", async () => {