fix tokenAmountNeededAndSTETHPayout to take into account stETH allowance

nounsDAO · Dec 30, 2024 · 92df5a9 · 92df5a9
1 parent 10b3660
commit 92df5a9
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 19 deletions.
diff --git a/src/STETHTokenBuyer.sol b/src/STETHTokenBuyer.sol
@@ -302,24 +302,24 @@ contract STETHTokenBuyer is Ownable, Pausable, ReentrancyGuard {
     }
 
     /// @notice Returns the amount of tokens the contract can buy and the amount of STETH it will pay for it
-    /// This takes into account the current STETH balance this contract has
+    /// This takes into account the current STETH allowance this contract has
     /// @return tokenAmount amount of tokens the contract can buy
-    /// @return ethAmount amount of STETH it will pay for the tokens
+    /// @return stethAmount amount of STETH it will pay for the tokens
     function tokenAmountNeededAndSTETHPayout() public view returns (uint256, uint256) {
         uint256 tokenAmount = tokenAmountNeeded();
-        uint256 ethAmount = stethAmountPerTokenAmount(tokenAmount);
-        uint256 ethAvailable = address(this).balance;
+        uint256 stethAmount = stethAmountPerTokenAmount(tokenAmount);
+        uint256 stethAvailable = stETH.allowance(treasury, address(this));
 
-        if (ethAvailable >= ethAmount) {
-            return (tokenAmount, ethAmount);
+        if (stethAvailable >= stethAmount) {
+            return (tokenAmount, stethAmount);
         } else {
             // Tokens amount will be rounded down to avoid trying to buy more eth than available
-            tokenAmount = tokenAmountPerSTEthAmount(ethAvailable);
+            tokenAmount = tokenAmountPerSTEthAmount(stethAvailable);
 
             // Recalculate eth amount because tokens amount are rounded down
-            ethAmount = stethAmountPerTokenAmount(tokenAmount);
+            stethAmount = stethAmountPerTokenAmount(tokenAmount);
 
-            return (tokenAmount, ethAmount);
+            return (tokenAmount, stethAmount);
         }
     }
 

diff --git a/test/STETHTokenBuyer.t.sol b/test/STETHTokenBuyer.t.sol
@@ -204,7 +204,8 @@ contract STETHTokenBuyerTest is Test {
     }
 
     function test_tokenAmountNeededAndETHPayout_baselineAmountOnly() public {
-        vm.deal(address(buyer), 50 ether);
+        vm.prank(treasury);
+        stETH.approve(address(buyer), 50 ether);
 
         vm.prank(owner);
         buyer.setBaselinePaymentTokenAmount(100_000e18);
@@ -231,14 +232,15 @@ contract STETHTokenBuyerTest is Test {
             _owner: owner,
             _admin: admin,
             _payer: address(payer),
-            _stETH: address(0),
-            _treasury: address(0)
+            _stETH: address(stETH),
+            _treasury: treasury
         });
 
         vm.prank(owner);
         buyer.setBaselinePaymentTokenAmount(100_000e6);
 
-        vm.deal(address(buyer), 8 ether);
+        vm.prank(treasury);
+        stETH.approve(address(buyer), 8 ether);
 
         priceFeed.setPrice(1350717518812290000000);
         (uint256 tokenAmount, uint256 ethAmount) = buyer.tokenAmountNeededAndSTETHPayout();
@@ -249,13 +251,13 @@ contract STETHTokenBuyerTest is Test {
     }
 
     function test_tokenAmountNeededAndETHPayout_lowersTokensIfItBuysMoreEthThanAvailable_fuzz(
-        uint256 ethBalance,
+        uint256 stethAllowance,
         uint256 price,
         uint256 decimals,
         uint256 tokensNeeded
     ) public {
         decimals = bound(decimals, 6, 18);
-        ethBalance = bound(ethBalance, 0, 1e12 ether);
+        stethAllowance = bound(stethAllowance, 0, 1e12 ether);
         price = bound(price, 1e18, 1e9 * 1e18);
         tokensNeeded = bound(tokensNeeded, 0, (10_000_000 * 10) ^ decimals);
 
@@ -273,14 +275,15 @@ contract STETHTokenBuyerTest is Test {
             _owner: owner,
             _admin: admin,
             _payer: address(payer),
-            _stETH: address(0),
-            _treasury: address(0)
+            _stETH: address(stETH),
+            _treasury: treasury
         });
 
         vm.prank(owner);
         buyer.setBaselinePaymentTokenAmount(tokensNeeded);
 
-        vm.deal(address(buyer), ethBalance);
+        vm.prank(treasury);
+        stETH.approve(address(buyer), stethAllowance);
 
         priceFeed.setPrice(price);
         (uint256 tokenAmount, uint256 ethAmount) = buyer.tokenAmountNeededAndSTETHPayout();
@@ -291,7 +294,8 @@ contract STETHTokenBuyerTest is Test {
     }
 
     function test_tokenAmountNeededAndETHPayout_lessEthAvailable() public {
-        vm.deal(address(buyer), 5 ether);
+        vm.prank(treasury);
+        stETH.approve(address(buyer), 5 ether);
 
         vm.prank(owner);
         buyer.setBaselinePaymentTokenAmount(100_000e18);