refactor: withdraw tests

morpho-org · Aug 4, 2023 · 2b17afa · 2b17afa
1 parent cb6e845
commit 2b17afa
Showing 1 changed file with 37 additions and 28 deletions.
diff --git a/test/forge/integration/TestIntegrationWithdraw.sol b/test/forge/integration/TestIntegrationWithdraw.sol
@@ -57,28 +57,34 @@ contract IntegrationWithdrawTest is BlueBaseTest {
         blue.withdraw(market, amountSupplied, address(this), address(this));
     }
 
-    function testWithdraw(uint256 amountSupplied, uint256 amountBorrowed, uint256 amountWithdrawn) public {
-        amountSupplied = bound(amountSupplied, 1, 2 ** 64);
-        amountBorrowed = bound(amountBorrowed, 1, amountSupplied);
-        amountWithdrawn = bound(amountWithdrawn, 1, amountSupplied);
-        vm.assume(amountWithdrawn <= amountSupplied - amountBorrowed);
+    function testWithdraw(
+        uint256 amountSupplied,
+        uint256 amountBorrowed,
+        uint256 amountWithdrawn,
+        address receiver
+    ) public {
+        vm.assume(receiver != address(0) && receiver != address(blue));
+
+        amountSupplied = bound(amountSupplied, 2, 2 ** 64);
+        amountBorrowed = bound(amountBorrowed, 1, amountSupplied - 1);
+        amountWithdrawn = bound(amountWithdrawn, 1, amountSupplied - amountBorrowed);
 
         borrowableAsset.setBalance(address(this), amountSupplied);
         blue.supply(market, amountSupplied, address(this), hex"");
 
         vm.prank(BORROWER);
         blue.borrow(market, amountBorrowed, BORROWER, BORROWER);
 
-        blue.withdraw(market, amountWithdrawn, address(this), address(this));
+        blue.withdraw(market, amountWithdrawn, address(this), receiver);
 
         assertApproxEqAbs(
             blue.supplyShares(id, address(this)),
             (amountSupplied - amountWithdrawn) * SharesMath.VIRTUAL_SHARES,
             100,
             "supply shares"
         );
-        assertEq(borrowableAsset.balanceOf(address(this)), amountWithdrawn, "this balance");
-        assertEq(borrowableAsset.balanceOf(BORROWER), amountBorrowed, "Borrower balance");
+        assertEq(borrowableAsset.balanceOf(receiver), amountWithdrawn, "receiver balance");
+        assertEq(borrowableAsset.balanceOf(BORROWER), amountBorrowed, "borrower balance");
         assertEq(
             borrowableAsset.balanceOf(address(blue)), amountSupplied - amountBorrowed - amountWithdrawn, "blue balance"
         );
@@ -88,38 +94,41 @@ contract IntegrationWithdrawTest is BlueBaseTest {
         uint256 amountSupplied,
         uint256 amountBorrowed,
         uint256 amountWithdrawn,
-        uint256 amountWithdrawnToReceiver
+        address onBehalf,
+        address receiver
     ) public {
-        amountSupplied = bound(amountSupplied, 1, 2 ** 64);
-        amountBorrowed = bound(amountBorrowed, 1, amountSupplied);
-        amountWithdrawn = bound(amountWithdrawn, 1, amountSupplied);
-        amountWithdrawnToReceiver = bound(amountWithdrawnToReceiver, 1, amountSupplied);
-        vm.assume(amountWithdrawn + amountWithdrawnToReceiver <= amountSupplied - amountBorrowed);
+        vm.assume(onBehalf != address(0) && onBehalf != address(blue));
+        vm.assume(receiver != address(0) && receiver != address(blue));
 
-        borrowableAsset.setBalance(address(this), amountSupplied);
-        blue.supply(market, amountSupplied, address(this), hex"");
-        blue.setAuthorization(BORROWER, true);
+        amountSupplied = bound(amountSupplied, 2, 2 ** 64);
+        amountBorrowed = bound(amountBorrowed, 1, amountSupplied - 1);
+        amountWithdrawn = bound(amountWithdrawn, 1, amountSupplied - amountBorrowed);
 
-        vm.startPrank(BORROWER);
-        blue.borrow(market, amountBorrowed, BORROWER, BORROWER);
-
-        blue.withdraw(market, amountWithdrawn, address(this), BORROWER);
-        blue.withdraw(market, amountWithdrawnToReceiver, address(this), address(this));
+        borrowableAsset.setBalance(onBehalf, amountSupplied);
 
+        vm.startPrank(onBehalf);
+        borrowableAsset.approve(address(blue), amountSupplied);
+        blue.supply(market, amountSupplied, onBehalf, hex"");
+        blue.borrow(market, amountBorrowed, onBehalf, onBehalf);
+        blue.setAuthorization(BORROWER, true);
         vm.stopPrank();
 
-        assertEq(blue.totalSupply(id), amountSupplied - amountWithdrawn - amountWithdrawnToReceiver, "total supply");
+        uint256 receiverBalanceBefore = borrowableAsset.balanceOf(receiver);
+
+        vm.startPrank(BORROWER);
+        blue.withdraw(market, amountWithdrawn, onBehalf, receiver);
+
+        assertEq(blue.totalSupply(id), amountSupplied - amountWithdrawn, "total supply");
         assertApproxEqAbs(
-            blue.supplyShares(id, address(this)),
-            (amountSupplied - amountWithdrawn - amountWithdrawnToReceiver) * SharesMath.VIRTUAL_SHARES,
+            blue.supplyShares(id, onBehalf),
+            (amountSupplied - amountWithdrawn) * SharesMath.VIRTUAL_SHARES,
             100,
             "supply shares"
         );
-        assertEq(borrowableAsset.balanceOf(address(this)), amountWithdrawnToReceiver, "this balance");
-        assertEq(borrowableAsset.balanceOf(BORROWER), amountBorrowed + amountWithdrawn, "Borrower balance");
+        assertEq(borrowableAsset.balanceOf(receiver) - receiverBalanceBefore, amountWithdrawn, "receiver balance");
         assertEq(
             borrowableAsset.balanceOf(address(blue)),
-            amountSupplied - amountBorrowed - amountWithdrawn - amountWithdrawnToReceiver,
+            amountSupplied - amountBorrowed - amountWithdrawn,
             "blue balance"
         );
     }