From 8348fe510fddbe681a1ea055ec65f08c310a0e1c Mon Sep 17 00:00:00 2001
From: smol-ninja <shubhamy2015@gmail.com>
Date: Wed, 13 Nov 2024 14:38:21 +0000
Subject: [PATCH] feat: revert depletionTimeOf if balance is zero (#335)

* feat: revert depletionTimeOf if balance is zero

* test: fix function name
---
 src/SablierFlow.sol                                       | 4 ++--
 src/interfaces/ISablierFlow.sol                           | 4 +++-
 src/libraries/Errors.sol                                  | 3 +++
 .../concrete/depletion-time-of/depletionTimeOf.t.sol      | 8 ++++----
 .../concrete/depletion-time-of/depletionTimeOf.tree       | 2 +-
 5 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/SablierFlow.sol b/src/SablierFlow.sol
index 5b0487d6..ea5c7cd1 100644
--- a/src/SablierFlow.sol
+++ b/src/SablierFlow.sol
@@ -64,9 +64,9 @@ contract SablierFlow is
     {
         uint128 balance = _streams[streamId].balance;
 
-        // If the stream balance is zero, return zero.
+        // If the stream balance is zero, revert to avoid ambiguity from a depleting stream.
         if (balance == 0) {
-            return 0;
+            revert Errors.SablierFlow_StreamBalanceZero(streamId);
         }
 
         uint8 tokenDecimals = _streams[streamId].tokenDecimals;
diff --git a/src/interfaces/ISablierFlow.sol b/src/interfaces/ISablierFlow.sol
index 41feb991..bb2c5a08 100644
--- a/src/interfaces/ISablierFlow.sol
+++ b/src/interfaces/ISablierFlow.sol
@@ -119,7 +119,9 @@ interface ISablierFlow is
 
     /// @notice Returns the time at which the total debt exceeds stream balance. If the total debt is less than
     /// or equal to stream balance, it returns 0.
-    /// @dev Reverts if `streamId` references a paused or a null stream.
+    /// @dev Reverts on the following conditions:
+    /// - If `streamId` references a paused or a null stream.
+    /// - If stream balance is zero.
     /// @param streamId The stream ID for the query.
     function depletionTimeOf(uint256 streamId) external view returns (uint256 depletionTime);
 
diff --git a/src/libraries/Errors.sol b/src/libraries/Errors.sol
index 28a2b4c2..080bcd3c 100644
--- a/src/libraries/Errors.sol
+++ b/src/libraries/Errors.sol
@@ -63,6 +63,9 @@ library Errors {
     /// @notice Thrown when trying to create a stream with the sender as the zero address.
     error SablierFlow_SenderZeroAddress();
 
+    /// @notice Thrown when trying to get depletion time of a stream with zero balance.
+    error SablierFlow_StreamBalanceZero(uint256 streamId);
+
     /// @notice Thrown when trying to perform an action with a paused stream.
     error SablierFlow_StreamPaused(uint256 streamId);
 
diff --git a/tests/integration/concrete/depletion-time-of/depletionTimeOf.t.sol b/tests/integration/concrete/depletion-time-of/depletionTimeOf.t.sol
index 8b52f10a..aca40c87 100644
--- a/tests/integration/concrete/depletion-time-of/depletionTimeOf.t.sol
+++ b/tests/integration/concrete/depletion-time-of/depletionTimeOf.t.sol
@@ -2,6 +2,7 @@
 pragma solidity >=0.8.22;
 
 import { UD21x18 } from "@prb/math/src/UD21x18.sol";
+import { Errors } from "src/libraries/Errors.sol";
 
 import { Integration_Test } from "../../Integration.t.sol";
 
@@ -16,10 +17,9 @@ contract DepletionTimeOf_Integration_Concrete_Test is Integration_Test {
         expectRevert_Paused(callData);
     }
 
-    function test_GivenBalanceZero() external view givenNotNull givenNotPaused {
-        // It should return 0.
-        uint256 actualDepletionTime = flow.depletionTimeOf(defaultStreamId);
-        assertEq(actualDepletionTime, 0, "depletion time");
+    function test_RevertGiven_BalanceZero() external givenNotNull givenNotPaused {
+        vm.expectRevert(abi.encodeWithSelector(Errors.SablierFlow_StreamBalanceZero.selector, defaultStreamId));
+        flow.depletionTimeOf(defaultStreamId);
     }
 
     function test_GivenUncoveredDebt() external givenNotNull givenNotPaused givenBalanceNotZero {
diff --git a/tests/integration/concrete/depletion-time-of/depletionTimeOf.tree b/tests/integration/concrete/depletion-time-of/depletionTimeOf.tree
index 3efa2717..bc3bbb57 100644
--- a/tests/integration/concrete/depletion-time-of/depletionTimeOf.tree
+++ b/tests/integration/concrete/depletion-time-of/depletionTimeOf.tree
@@ -6,7 +6,7 @@ DepletionTimeOf_Integration_Concrete_Test
    │  └── it should revert
    └── given not paused
       ├── given balance zero
-      │  └── it should return 0
+      │  └── it should revert
       └── given balance not zero
          ├── given uncovered debt
          │  └── it should return 0