account for total pending in availableToWithdraw.

reserve-protocol · Sep 27, 2023 · b3c2bb8 · b3c2bb8
1 parent 0e7f635
commit b3c2bb8
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/src/ThrottleWallet.sol b/src/ThrottleWallet.sol
@@ -155,7 +155,7 @@ contract ThrottleWallet {
  accumulatedWithdrawalAmount = amountPerPeriod;
  }
 
- uint256 availableBalance = throttledToken.balanceOf(address(this));
+ uint256 availableBalance = throttledToken.balanceOf(address(this)) - totalPending;
  if (accumulatedWithdrawalAmount > availableBalance) {
  accumulatedWithdrawalAmount = availableBalance;
  }

diff --git a/test/ThrottleWallet.t.sol b/test/ThrottleWallet.t.sol
@@ -172,12 +172,15 @@ contract ThrottleWalletTest is Test {
  }
 
  function test_LinearAccumulator() public {
+ // start w/ 4B tokens
+ token.mint(address(throttleWallet), 2_000_000_000 ether);
+
  vm.startPrank(user_user);
 
  // Drain the entire throttle!
  throttleWallet.initiateWithdrawal(1_000_000_000 ether, user_target);
 
- assertEq(token.balanceOf(address(throttleWallet)), 2_000_000_000 ether);
+ assertEq(token.balanceOf(address(throttleWallet)), 4_000_000_000 ether);
  assertEq(token.balanceOf(address(this)), 0);
  assertEq(throttleWallet.totalPending(), 1_000_000_000 ether);
 
@@ -189,7 +192,7 @@ contract ThrottleWalletTest is Test {
  vm.warp(START_TIME + 2 weeks);
  throttleWallet.initiateWithdrawal(500_000_000 ether, user_target);
 
- assertEq(token.balanceOf(address(throttleWallet)), 2_000_000_000 ether);
+ assertEq(token.balanceOf(address(throttleWallet)), 4_000_000_000 ether);
  assertEq(token.balanceOf(address(this)), 0);
  assertEq(throttleWallet.totalPending(), 1_500_000_000 ether);
 
@@ -210,7 +213,7 @@ contract ThrottleWalletTest is Test {
  throttleWallet.completeWithdrawal(1);
 
  // The throttle is now full charged
- assertEq(throttleWallet.availableToWithdraw(), 500_000_000 ether);
+ assertEq(throttleWallet.availableToWithdraw(), 1_000_000_000 ether);
  vm.stopPrank();
  }
 
@@ -257,6 +260,12 @@ contract ThrottleWalletTest is Test {
  throttleWallet.completeWithdrawal(0);
 
  throttleWallet.initiateWithdrawal(500_000_000 ether, user_target);
+
+ // total pending is accounted for
+ assertEq(throttleWallet.availableToWithdraw(), 500_000_000 ether);
+ vm.warp(START_TIME + 5 weeks);
+ assertEq(throttleWallet.availableToWithdraw(), 500_000_000 ether);
+
  vm.warp(START_TIME + 8 weeks);
  throttleWallet.completeWithdrawal(1);