test: add yieldWithdraw

value4value · May 21, 2024 · 28e8ec0 · 28e8ec0
1 parent 3be73a8
commit 28e8ec0
Showing 1 changed file with 41 additions and 51 deletions.
diff --git a/test/unit/YieldAggregator.t.sol b/test/unit/YieldAggregator.t.sol
@@ -2,7 +2,6 @@
 
 pragma solidity 0.8.25;
 
-import { console } from "forge-std/console.sol";
 import { BaseTest } from "../BaseTest.t.sol";
 
 contract YieldAggregatorTests is BaseTest {
@@ -12,15 +11,6 @@ contract YieldAggregatorTests is BaseTest {
         createFactory();
     }
 
-    function test_migrateNewYieldAggregator() public {
-        vm.prank(owner);
-        vm.expectRevert(bytes("Invalid yieldAggregator"));
-        sharesFactory.migrate(address(0));
-
-        vm.prank(owner);
-        sharesFactory.migrate(address(aaveYieldAggregator));
-    }
-
     function test_setYieldBuffer() public {
         vm.prank(owner);
         aaveYieldAggregator.setYieldBuffer(1e11);
@@ -32,17 +22,11 @@ contract YieldAggregatorTests is BaseTest {
     }
 
     function test_yieldDepost() public {
-        // Mock eth balance, and call yieldDeposit
-        vm.deal(address(aaveYieldAggregator), 10 ether);
-        vm.prank(address(sharesFactory));
-        aaveYieldAggregator.yieldDeposit();
+        _depositETH2AaveYieldAggregator(10 ether);
         assertTrue(aWETH.balanceOf(address(sharesFactory)) == 10 ether);
         assertTrue(address(aaveYieldAggregator).balance == 0);
 
-        // Migrate to BlankYieldAggregator
-        // Mock eth balance, and call yieldDeposit
-        vm.prank(owner);
-        sharesFactory.migrate(address(blankYieldAggregator));
+        _migrate2BlankYieldAggregator();
         vm.deal(address(blankYieldAggregator), 10 ether);
         vm.prank(address(sharesFactory));
         blankYieldAggregator.yieldDeposit();
@@ -51,56 +35,62 @@ contract YieldAggregatorTests is BaseTest {
 
         vm.expectRevert(bytes("Only factory"));
         aaveYieldAggregator.yieldDeposit();
+
+        vm.expectRevert(bytes("Only factory"));
+        blankYieldAggregator.yieldDeposit();
     }
 
     function test_yieldWithdraw() public {
-        // Sell shares
-        // Claim yield
-        // Migrate YieldAggregator
-        // Check aWETH / ETH balance
+        _depositETH2AaveYieldAggregator(10 ether);
+
+        vm.prank(address(sharesFactory));
+        aaveYieldAggregator.yieldWithdraw(10 ether);
+        assertTrue(aWETH.balanceOf(address(sharesFactory)) == 0);
+        assertTrue(address(sharesFactory).balance == 10 ether);
+
+        _migrate2BlankYieldAggregator();
+        vm.prank(address(sharesFactory));
+        blankYieldAggregator.yieldWithdraw(10 ether);
+        assertTrue(address(blankYieldAggregator).balance == 0);
+        assertTrue(address(sharesFactory).balance == 10 ether);
+
+        vm.expectRevert(bytes("Only factory"));
+        aaveYieldAggregator.yieldWithdraw(10 ether);
+
+        vm.expectRevert(bytes("Only factory"));
+        blankYieldAggregator.yieldWithdraw(0);
     }
 
     function test_yieldBalanceOf() public {
+        _depositETH2AaveYieldAggregator(10 ether);
+
+        uint256 aaveYieldBalance = aaveYieldAggregator.yieldBalanceOf(address(sharesFactory));
+        assertEq(aaveYieldBalance, 10 ether);
 
+        uint256 blankYieldBalance = blankYieldAggregator.yieldBalanceOf(address(sharesFactory));
+        assertEq(blankYieldBalance, 0);
     }
 
-    function test_yieldToken() public {
+    function test_yieldToken() public view {
+        address aaveYieldToken = aaveYieldAggregator.yieldToken();
+        assertEq(aaveYieldToken, address(aWETH));
 
+        address blankYieldToken = blankYieldAggregator.yieldToken();
+        assertEq(blankYieldToken, address(WETH));
     }
 
     function test_yieldMaxClaimable() public {
 
     }
 
-    function _testBuyShares() public {
-        uint256 aliceBalBefore = addrAlice.balance;
-        uint256 bobBalBefore = addrBob.balance;
-        uint256 referrerBalBefore = referralReceiver.balance;
-        // uint256 factoryBalBefore = aWETH.balanceOf(address(sharesFactory));
-        uint256 depositedETHAmountBefore = sharesFactory.depositedETHAmount();
-
-        vm.deal(addrBob, 10 ether);
-        _buyShare(addrBob, 0, 2, referralReceiver);
-
-        uint256 aliceBalAfter = addrAlice.balance;
-        uint256 bobBalAfter = addrBob.balance;
-        uint256 referrerBalAfter = referralReceiver.balance;
-        // uint256 factoryBalAfter = aWETH.balanceOf(address(sharesFactory));
-        uint256 depositedETHAmountAfter = sharesFactory.depositedETHAmount();
-
-        assertEq(bobBalBefore - bobBalAfter, 5500450999999993); // Bob buy 1 share
-        assertEq(aliceBalAfter - aliceBalBefore, 250020499999999); // Alice receive creator fee
-        assertEq(referrerBalAfter - referrerBalBefore, 250020499999999); // referral receive fee
-        // assertEq(factoryBalAfter - factoryBalBefore, 5000409999999995); // Factory aWETH balance with rounding error
-        assertEq(depositedETHAmountAfter - depositedETHAmountBefore, 5000409999999995); // Factory records ETH Amount
-
-        uint256 bobShareBal = sharesNFT.balanceOf(addrBob, 0);
-        assertEq(bobShareBal, 2);
+    function _depositETH2AaveYieldAggregator(uint256 amount) internal {
+        vm.deal(address(aaveYieldAggregator), amount);
+        vm.prank(address(sharesFactory));
+        aaveYieldAggregator.yieldDeposit();
     }
 
-    function _buyShare(address sender, uint256 shareId, uint32 quantity, address referral) internal {
-        (uint256 buyPriceAfterFee,,,) = sharesFactory.getBuyPriceAfterFee(shareId, quantity, referral);
-        vm.prank(address(sender));
-        sharesFactory.buyShare{ value: buyPriceAfterFee }(shareId, quantity, referral);
+    function _migrate2BlankYieldAggregator() internal {
+        vm.prank(owner);
+        sharesFactory.migrate(address(blankYieldAggregator));
     }
 }