Merge pull request #10 from metadock/tests/withdraw-stream

Implement `withdrawLinearStream` and `withdrawTranchedStream` integration tests

src/modules/invoice-module/sablier-v2/StreamManager.sol

@@ -125,13 +125,13 @@ abstract contract StreamManager is IStreamManager {
     //////////////////////////////////////////////////////////////////////////*/
 
     /// @inheritdoc IStreamManager
-    function withdrawLinearStream(uint256 streamId, address to, uint128 amount) public {
-        _withdrawStream({ sablier: LOCKUP_LINEAR, streamId: streamId, to: to, amount: amount });
+    function withdrawLinearStream(uint256 streamId, address to) public {
+        _withdrawStream({ sablier: LOCKUP_LINEAR, streamId: streamId, to: to });
     }
 
     /// @inheritdoc IStreamManager
-    function withdrawTranchedStream(uint256 streamId, address to, uint128 amount) public {
-        _withdrawStream({ sablier: LOCKUP_TRANCHED, streamId: streamId, to: to, amount: amount });
+    function withdrawTranchedStream(uint256 streamId, address to) public {
+        _withdrawStream({ sablier: LOCKUP_TRANCHED, streamId: streamId, to: to });
     }
 
     /*//////////////////////////////////////////////////////////////////////////
@@ -244,8 +244,8 @@ abstract contract StreamManager is IStreamManager {
     }
 
     /// @dev Withdraws from either a linear or tranched stream
-    function _withdrawStream(ISablierV2Lockup sablier, uint256 streamId, address to, uint128 amount) internal {
-        sablier.withdraw(streamId, to, amount);
+    function _withdrawStream(ISablierV2Lockup sablier, uint256 streamId, address to) internal {
+        sablier.withdrawMax(streamId, to);
     }
 
     /// @dev Cancels the `streamId` stream

src/modules/invoice-module/sablier-v2/interfaces/IStreamManager.sol

@@ -93,11 +93,11 @@ interface IStreamManager {
     /// @param newBrokerFee The new broker fee
     function updateStreamBrokerFee(UD60x18 newBrokerFee) external;
 
-    /// @notice See the documentation in {ISablierV2Lockup-withdraw}
-    function withdrawLinearStream(uint256 streamId, address to, uint128 amount) external;
+    /// @notice See the documentation in {ISablierV2Lockup-withdrawMax}
+    function withdrawLinearStream(uint256 streamId, address to) external;
 
-    /// @notice See the documentation in {ISablierV2Lockup-withdraw}
-    function withdrawTranchedStream(uint256 streamId, address to, uint128 amount) external;
+    /// @notice See the documentation in {ISablierV2Lockup-withdrawMax}
+    function withdrawTranchedStream(uint256 streamId, address to) external;
 
     /// @notice See the documentation in {ISablierV2Lockup-cancel}
     ///

test/integration/concrete/invoice-module/withdraw-linear-stream/withdrawLinearStream.t.sol

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { WithdrawLinearStream_Integration_Shared_Test } from "../../../shared/withdrawLinearStream.t.sol";
+
+contract WithdrawLinearStream_Integration_Concret_Test is WithdrawLinearStream_Integration_Shared_Test {
+    function setUp() public virtual override {
+        WithdrawLinearStream_Integration_Shared_Test.setUp();
+    }
+
+    function test_WithdrawLinearStream() external givenPaymentMethodLinearStream givenInvoiceStatusOngoing {
+        // Set current invoice as a linear stream-based one
+        uint256 invoiceId = 3;
+        uint256 streamId = 1;
+
+        // The invoice must be paid for its status to be updated to `Ongoing`
+        // Make Bob the payer of the invoice (also Bob will be the initial stream sender)
+        vm.startPrank({ msgSender: users.bob });
+
+        // Approve the {InvoiceModule} to transfer the USDT tokens on Bob's behalf
+        usdt.approve({ spender: address(invoiceModule), amount: invoices[invoiceId].payment.amount });
+
+        // Pay the invoice first (status will be updated to `Ongoing`)
+        invoiceModule.payInvoice{ value: invoices[invoiceId].payment.amount }({ id: invoiceId });
+
+        // Advance the timestamp by 3 weeks to simulate the withdrawal
+        vm.warp(block.timestamp + 3 weeks);
+
+        // Store Eve's balance before withdrawing the USDT tokens
+        uint256 balanceOfBefore = usdt.balanceOf(users.eve);
+
+        // Get the maximum withdrawable amount from the stream
+        uint128 maxWithdrawableAmount = sablierV2LockupLinear.withdrawableAmountOf(streamId);
+
+        // Make Eve the caller in this test suite as she's the recipient of the invoice
+        vm.startPrank({ msgSender: users.eve });
+
+        // Run the test
+        invoiceModule.withdrawLinearStream({ streamId: streamId, to: users.eve });
+
+        // Assert the current and expected USDT balance of Eve
+        assertEq(balanceOfBefore + maxWithdrawableAmount, usdt.balanceOf(users.eve));
+    }
+}

test/integration/concrete/invoice-module/withdraw-linear-stream/withdrawLinearStream.tree

@@ -0,0 +1,4 @@
+withdrawLinearStream.t.sol
+└── given the payment method is linear stream
+    └── given the invoice status is Ongoing
+        └── it should allow the invoice recipient to withdraw from the stream

test/integration/concrete/invoice-module/withdraw-tranched-stream/withdrawTranchedStream.t.sol

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { WithdrawTranchedStream_Integration_Shared_Test } from "../../../shared/withdrawTranchedStream.t.sol";
+
+contract WithdrawTranchedStream_Integration_Concret_Test is WithdrawTranchedStream_Integration_Shared_Test {
+    function setUp() public virtual override {
+        WithdrawTranchedStream_Integration_Shared_Test.setUp();
+    }
+
+    function test_WithdrawTranchedStream() external givenPaymentMethodTranchedStream givenInvoiceStatusOngoing {
+        // Set current invoice as a tranched stream-based one
+        uint256 invoiceId = 4;
+        uint256 streamId = 1;
+
+        // The invoice must be paid for its status to be updated to `Ongoing`
+        // Make Bob the payer of the invoice (also Bob will be the initial stream sender)
+        vm.startPrank({ msgSender: users.bob });
+
+        // Approve the {InvoiceModule} to transfer the USDT tokens on Bob's behalf
+        usdt.approve({ spender: address(invoiceModule), amount: invoices[invoiceId].payment.amount });
+
+        // Pay the invoice first (status will be updated to `Ongoing`)
+        invoiceModule.payInvoice{ value: invoices[invoiceId].payment.amount }({ id: invoiceId });
+
+        // Advance the timestamp by 3 weeks to simulate the withdrawal
+        vm.warp(block.timestamp + 3 weeks);
+
+        // Store Eve's balance before withdrawing the USDT tokens
+        uint256 balanceOfBefore = usdt.balanceOf(users.eve);
+
+        // Get the maximum withdrawable amount from the stream
+        uint128 maxWithdrawableAmount = sablierV2LockupTranched.withdrawableAmountOf(streamId);
+
+        // Make Eve the caller in this test suite as she's the recipient of the invoice
+        vm.startPrank({ msgSender: users.eve });
+
+        // Run the test
+        invoiceModule.withdrawTranchedStream({ streamId: streamId, to: users.eve });
+
+        // Assert the current and expected USDT balance of Eve
+        assertEq(balanceOfBefore + maxWithdrawableAmount, usdt.balanceOf(users.eve));
+    }
+}

test/integration/concrete/invoice-module/withdraw-tranched-stream/withdrawTranchedStream.tree

@@ -0,0 +1,4 @@
+withdrawTranchedStream.t.sol
+└── given the payment method is tranched stream
+    └── given the invoice status is Ongoing
+        └── it should allow the invoice recipient to withdraw from the stream

test/integration/shared/withdrawLinearStream.t.sol

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { Integration_Test } from "../Integration.t.sol";
+import { PayInvoice_Integration_Shared_Test } from "./payInvoice.t.sol";
+
+abstract contract WithdrawLinearStream_Integration_Shared_Test is Integration_Test, PayInvoice_Integration_Shared_Test {
+    function setUp() public virtual override(Integration_Test, PayInvoice_Integration_Shared_Test) {
+        PayInvoice_Integration_Shared_Test.setUp();
+    }
+
+    modifier givenInvoiceStatusOngoing() {
+        _;
+    }
+}

test/integration/shared/withdrawTranchedStream.t.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { Integration_Test } from "../Integration.t.sol";
+import { PayInvoice_Integration_Shared_Test } from "./payInvoice.t.sol";
+
+abstract contract WithdrawTranchedStream_Integration_Shared_Test is
+    Integration_Test,
+    PayInvoice_Integration_Shared_Test
+{
+    function setUp() public virtual override(Integration_Test, PayInvoice_Integration_Shared_Test) {
+        PayInvoice_Integration_Shared_Test.setUp();
+    }
+
+    modifier givenInvoiceStatusOngoing() {
+        _;
+    }
+}