From e2ec3bea0455ebb42e3b5be5d927c5c11f1a87bf Mon Sep 17 00:00:00 2001
From: Marc Doerflinger <marc@doerf.li>
Date: Tue, 9 Jul 2024 08:07:44 +0000
Subject: [PATCH] add unit test (#476)

---
 contracts/pool/PoolService.sol | 28 +++++++++++++----------
 test/TestBundle.t.sol          | 41 ++++++++++++++++++++++++++++++++++
 test/TestPool.t.sol            |  4 ++--
 3 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/contracts/pool/PoolService.sol b/contracts/pool/PoolService.sol
index f777a480d..b1f496812 100644
--- a/contracts/pool/PoolService.sol
+++ b/contracts/pool/PoolService.sol
@@ -118,21 +118,23 @@ contract PoolService is
         returns(NftId bundleNftId, Amount netStakedAmount)
     {
         (NftId poolNftId,, IInstance instance) = _getAndVerifyActiveComponent(POOL());
-        
-        Amount stakingFeeAmount;
-        (stakingFeeAmount, netStakedAmount) = FeeLib.calculateFee(
-            _getStakingFee(instance.getInstanceReader(), poolNftId), 
-            stakingAmount);
 
         {
             InstanceReader instanceReader = instance.getInstanceReader();
             IComponents.PoolInfo memory poolInfo = instanceReader.getPoolInfo(poolNftId);
             Amount currentPoolBalance = instanceReader.getBalanceAmount(poolNftId);
-            if (currentPoolBalance + netStakedAmount > poolInfo.maxBalanceAmount) {
-                revert ErrorPoolServiceMaxBalanceAmountExceeded(poolNftId, poolInfo.maxBalanceAmount, currentPoolBalance, netStakedAmount);
+            if (currentPoolBalance + stakingAmount > poolInfo.maxBalanceAmount) {
+                revert ErrorPoolServiceMaxBalanceAmountExceeded(poolNftId, poolInfo.maxBalanceAmount, currentPoolBalance, stakingAmount);
             }
         }
 
+
+        Amount stakingFeeAmount;
+        (stakingFeeAmount, netStakedAmount) = FeeLib.calculateFee(
+            _getStakingFee(instance.getInstanceReader(), poolNftId), 
+            stakingAmount);
+
+
         bundleNftId = _bundleService.create(
             instance,
             poolNftId,
@@ -201,6 +203,13 @@ contract PoolService is
             revert ErrorPoolServiceBundlePoolMismatch(bundleNftId, poolNftId);
         }
 
+        {
+            Amount currentPoolBalance = instanceReader.getBalanceAmount(poolNftId);
+            if (currentPoolBalance + amount > poolInfo.maxBalanceAmount) {
+                revert ErrorPoolServiceMaxBalanceAmountExceeded(poolNftId, poolInfo.maxBalanceAmount, currentPoolBalance, amount);
+            }
+        }
+
         // calculate fees
         Amount feeAmount;
         (
@@ -210,11 +219,6 @@ contract PoolService is
             _getStakingFee(instanceReader, poolNftId), 
             amount);
 
-        Amount currentPoolBalance = instanceReader.getBalanceAmount(poolNftId);
-        if (currentPoolBalance + netAmount > poolInfo.maxBalanceAmount) {
-            revert ErrorPoolServiceMaxBalanceAmountExceeded(poolNftId, poolInfo.maxBalanceAmount, currentPoolBalance, netAmount);
-        }
-
         // do all the bookkeeping
         _componentService.increasePoolBalance(
             instance.getInstanceStore(), 
diff --git a/test/TestBundle.t.sol b/test/TestBundle.t.sol
index bd03304a4..b8c8da2d3 100644
--- a/test/TestBundle.t.sol
+++ b/test/TestBundle.t.sol
@@ -135,6 +135,47 @@ contract TestBundle is GifTest {
         assertEq(instanceReader.getFeeAmount(bundleNftId).toInt(), 0, "bundle fees 0");
     }
 
+    function test_Bundle_stakeBundle_maxBalanceExceeded() public {
+        // GIVEN
+        initialStakingFee = FeeLib.percentageFee(4);
+        _prepareProduct(false);
+        
+        IComponents.ComponentInfo memory poolComponentInfo = instanceReader.getComponentInfo(poolNftId);
+        vm.startPrank(poolOwner);
+        pool.setMaxBalanceAmount(AmountLib.toAmount(1500));
+        vm.stopPrank();
+
+        vm.startPrank(investor);
+        token.approve(address(pool.getTokenHandler()), 2000);
+
+        Seconds lifetime = SecondsLib.toSeconds(604800);
+        (bundleNftId,) = pool.createBundle(
+            FeeLib.zero(), 
+            1000, 
+            lifetime, 
+            ""
+        );
+
+        assertTrue(!bundleNftId.eqz(), "bundle nft id is zero");
+
+        uint256 stakeAmount = 1001;
+        Amount stakeAmt = AmountLib.toAmount(stakeAmount);
+        Amount stakeNetAmt = AmountLib.toAmount(960);
+
+        pool.lockBundle(bundleNftId);
+
+        // THEN - expect revert
+        vm.expectRevert(abi.encodeWithSelector(
+            IPoolService.ErrorPoolServiceMaxBalanceAmountExceeded.selector, 
+            poolNftId,
+            AmountLib.toAmount(1500),
+            AmountLib.toAmount(1000),
+            AmountLib.toAmount(1001)));
+
+        // WHEN - pool is staked with more tokens
+        pool.stake(bundleNftId, stakeAmt);
+    }
+
     /// @dev test staking when the allowance is too small
     function test_Bundle_stakeBundle_allowanceTooSmall() public {
         // GIVEN
diff --git a/test/TestPool.t.sol b/test/TestPool.t.sol
index e8543651e..7e42e1176 100644
--- a/test/TestPool.t.sol
+++ b/test/TestPool.t.sol
@@ -251,8 +251,8 @@ contract TestPool is GifTest {
         vm.expectRevert(abi.encodeWithSelector(
             IPoolService.ErrorPoolServiceMaxBalanceAmountExceeded.selector, 
             poolNftId,
-            AmountLib.toAmount(10000),
-            AmountLib.toAmount(10000),
+            AmountLib.toAmount(15000),
+            AmountLib.toAmount(100000000000),
             AmountLib.toAmount(10000)));
 
         // WHEN