Remove amount arg in withdraw and withdraw all the withdrawable amount

smartcontractkit · Dec 18, 2023 · 94b1d53 · 94b1d53
1 parent 54a523e
commit 94b1d53
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/contracts/src/v0.8/vrf/dev/SubscriptionAPI.sol b/contracts/src/v0.8/vrf/dev/SubscriptionAPI.sol
@@ -208,13 +208,14 @@ abstract contract SubscriptionAPI is ConfirmedOwner, IERC677Receiver, IVRFSubscr
    * @param recipient where to send the funds
    * @param amount amount to withdraw
    */
-  function withdraw(address recipient, uint96 amount) external nonReentrant onlyOwner {
+  function withdraw(address recipient) external nonReentrant onlyOwner {
     if (address(LINK) == address(0)) {
       revert LinkNotSet();
     }
-    if (s_withdrawableTokens < amount) {
+    if (s_withdrawableTokens == 0) {
       revert InsufficientBalance();
     }
+    uint96 amount = s_withdrawableTokens;
     s_withdrawableTokens -= amount;
     s_totalBalance -= amount;
     if (!LINK.transfer(recipient, amount)) {
@@ -227,11 +228,12 @@ abstract contract SubscriptionAPI is ConfirmedOwner, IERC677Receiver, IVRFSubscr
    * @param recipient where to send the funds
    * @param amount amount to withdraw
    */
-  function withdrawNative(address payable recipient, uint96 amount) external nonReentrant onlyOwner {
-    if (s_withdrawableNative < amount) {
+  function withdrawNative(address payable recipient) external nonReentrant onlyOwner {
+    if (s_withdrawableNative == 0) {
       revert InsufficientBalance();
     }
     // Prevent re-entrancy by updating state before transfer.
+    uint96 amount = s_withdrawableNative;
     s_withdrawableNative -= amount;
     s_totalNativeBalance -= amount;
     (bool sent, ) = recipient.call{value: amount}("");

diff --git a/contracts/test/v0.8/foundry/vrf/VRFV2PlusSubscriptionAPI.t.sol b/contracts/test/v0.8/foundry/vrf/VRFV2PlusSubscriptionAPI.t.sol
@@ -318,7 +318,7 @@ contract VRFV2PlusSubscriptionAPITest is BaseTest {
   function testWithdrawNoLink() public {
     // CASE: no link token set
     vm.expectRevert(SubscriptionAPI.LinkNotSet.selector);
-    s_subscriptionAPI.withdraw(OWNER, 1 ether);
+    s_subscriptionAPI.withdraw(OWNER);
   }
 
   function testWithdrawInsufficientBalance() public {
@@ -330,7 +330,7 @@ contract VRFV2PlusSubscriptionAPITest is BaseTest {
 
     // call withdraw
     vm.expectRevert(SubscriptionAPI.InsufficientBalance.selector);
-    s_subscriptionAPI.withdraw(OWNER, 1 ether);
+    s_subscriptionAPI.withdraw(OWNER);
   }
 
   function testWithdrawSufficientBalanceLinkSet() public {
@@ -355,7 +355,7 @@ contract VRFV2PlusSubscriptionAPITest is BaseTest {
     // call Withdraw from owner address
     uint256 ownerBalance = linkToken.balanceOf(OWNER);
     changePrank(OWNER);
-    s_subscriptionAPI.withdraw(OWNER, 1 ether);
+    s_subscriptionAPI.withdraw(OWNER);
     // assert link balance of owner
     assertEq(linkToken.balanceOf(OWNER) - ownerBalance, 1 ether, "owner link balance incorrect");
     // assert state of subscription api
@@ -371,21 +371,21 @@ contract VRFV2PlusSubscriptionAPITest is BaseTest {
     // call WithdrawNative
     changePrank(OWNER);
     vm.expectRevert(SubscriptionAPI.InsufficientBalance.selector);
-    s_subscriptionAPI.withdrawNative(payable(OWNER), 1 ether);
+    s_subscriptionAPI.withdrawNative(payable(OWNER));
   }
 
   function testWithdrawLinkInvalidOwner() public {
     address invalidAddress = makeAddr("invalidAddress");
     changePrank(invalidAddress);
     vm.expectRevert("Only callable by owner");
-    s_subscriptionAPI.withdraw(payable(OWNER), 1 ether);
+    s_subscriptionAPI.withdraw(payable(OWNER));
   }
 
   function testWithdrawNativeInvalidOwner() public {
     address invalidAddress = makeAddr("invalidAddress");
     changePrank(invalidAddress);
     vm.expectRevert("Only callable by owner");
-    s_subscriptionAPI.withdrawNative(payable(OWNER), 1 ether);
+    s_subscriptionAPI.withdrawNative(payable(OWNER));
   }
 
   function testWithdrawNativeSufficientBalance() public {
@@ -405,7 +405,7 @@ contract VRFV2PlusSubscriptionAPITest is BaseTest {
 
     // call WithdrawNative from owner address
     changePrank(OWNER);
-    s_subscriptionAPI.withdrawNative(payable(OWNER), 1 ether);
+    s_subscriptionAPI.withdrawNative(payable(OWNER));
     // assert native balance
     assertEq(address(OWNER).balance, 1 ether, "owner native balance incorrect");
     // assert state of subscription api