update tests and fix admin fee claiming logic

contracts/main/CurveTricryptoOptimizedWETH.vy

@@ -1060,10 +1060,9 @@ def _claim_admin_fees():
     # 2. Pool parameters are being ramped.
 
     last_claim_time: uint256 = self.last_admin_fee_claim_timestamp
-
     if (
-        block.timestamp - last_claim_time < MIN_ADMIN_FEE_CLAIM_INTERVAL
-        or self.future_A_gamma_time < block.timestamp
+        block.timestamp - last_claim_time < MIN_ADMIN_FEE_CLAIM_INTERVAL or
+        self.future_A_gamma_time > block.timestamp
     ):
         return
 
@@ -1079,7 +1078,7 @@ def _claim_admin_fees():
     # 2. there are less than 10**18 (or 1 unit of) lp tokens, else it can lead
     #    to manipulated virtual prices.
 
-    if (xcp_profit <= xcp_profit_a or current_lp_token_supply < 10**18):
+    if xcp_profit <= xcp_profit_a or current_lp_token_supply < 10**18:
         return
 
     # ---------- Conditions met to claim admin fees: compute state. ----------
@@ -1203,12 +1202,13 @@ def _A_gamma() -> uint256[2]:
 def _fee(xp: uint256[N_COINS]) -> uint256:
 
     fee_params: uint256[3] = self._unpack(self.packed_fee_params)
+    f: uint256 = MATH.reduction_coefficient(xp, fee_params[2])
 
-    if self.future_A_gamma_time > block.timestamp:  # TODO: do not charge max fee else pool rekt!
-        fee_params[0] = MAX_FEE  # mid_fee is MAX_FEE during ramping
-        fee_params[1] = MAX_FEE  # out_fee is MAX_FEE during ramping
+    # During parameter ramping, we raise fees and disable admin fee claiming
+    if self.future_A_gamma_time > block.timestamp:  # parameter ramping
+        fee_params[0] = 10**8  # set mid_fee to 100 basis points
+        fee_params[1] = 10**8  # set out_fee to 100 basis points
 
-    f: uint256 = MATH.reduction_coefficient(xp, fee_params[2])
     return unsafe_div(
         fee_params[0] * f + fee_params[1] * (10**18 - f),
         10**18

tests/boa/unitary/pool/stateful/stateful_base.py

@@ -40,6 +40,7 @@ def __init__(self):
         self.user_balances = {u: [0] * 3 for u in self.accounts}
         self.balances = self.initial_deposit[:]
         self.xcp_profit = 10**18
+        self.xcp_profit_a = 10**18
 
         self.total_supply = 0
         self.previous_pool_profit = 0
@@ -259,11 +260,43 @@ def up_only_profit(self):
     @contextlib.contextmanager
     def upkeep_on_claim(self):
 
-        admin_balance = self.swap.balanceOf(self.fee_receiver)
+        admin_balances_pre = [
+            c.balanceOf(self.fee_receiver) for c in self.coins
+        ]
+        pool_is_ramping = (
+            self.swap.future_A_gamma_time() > boa.env.vm.state.timestamp
+        )
+
         try:
+
             yield
+
         finally:
-            _claimed = self.swap.balanceOf(self.fee_receiver) - admin_balance
-            if _claimed > 0:
-                self.total_supply += _claimed
-                self.xcp_profit = self.swap.xcp_profit()
+
+            new_xcp_profit_a = self.swap.xcp_profit_a()
+            old_xcp_profit_a = self.xcp_profit_a
+
+            claimed = False
+            if new_xcp_profit_a > old_xcp_profit_a:
+                claimed = True
+                self.xcp_profit_a = new_xcp_profit_a
+
+            admin_balances_post = [
+                c.balanceOf(self.fee_receiver) for c in self.coins
+            ]
+
+            if claimed:
+
+                for i in range(3):
+                    claimed_amount = (
+                        admin_balances_post[i] - admin_balances_pre[i]
+                    )
+                    assert (
+                        claimed_amount > 0
+                    )  # check if non zero amounts of claim
+                    assert not pool_is_ramping  # cannot claim while ramping
+
+                    # update self.balances
+                    self.balances[i] -= claimed_amount
+
+        self.xcp_profit = self.swap.xcp_profit()

tests/boa/unitary/pool/stateful/test_admin_fee_claim.py

@@ -3,7 +3,8 @@
 from boa.test import strategy
 from hypothesis.stateful import rule, run_state_machine_as_test
 
-from tests.boa.unitary.pool.stateful.stateful_base import StatefulBase
+# from tests.boa.unitary.pool.stateful.stateful_base import StatefulBase
+from tests.boa.unitary.pool.stateful.test_stateful import ProfitableState
 
 MAX_SAMPLES = 20
 STEP_COUNT = 100
@@ -16,10 +17,16 @@ def approx(x1, x2, precision):
     return abs(log(x1 / x2)) <= precision
 
 
-class StatefulAdmin(StatefulBase):
+class StatefulAdmin(ProfitableState):
     exchange_amount_in = strategy(
         "uint256", min_value=10**17, max_value=10**5 * 10**18
     )
+    deposit_amounts = strategy(
+        "uint256[3]", min_value=10**18, max_value=10**9 * 10**18
+    )
+    token_amount = strategy(
+        "uint256", min_value=10**18, max_value=10**12 * 10**18
+    )
     exchange_i = strategy("uint8", max_value=2)
     exchange_j = strategy("uint8", max_value=2)
     user = strategy("address")
@@ -54,15 +61,24 @@ def exchange(self, exchange_amount_in, exchange_i, exchange_j, user):
             exchange_amount_in_converted, exchange_i, exchange_j, user
         )
 
+    @rule(deposit_amounts=deposit_amounts, user=user)
+    def deposit(self, deposit_amounts, user):
+        deposit_amounts[1:] = [deposit_amounts[0]] + [
+            deposit_amounts[i] * 10**18 // self.swap.price_oracle(i - 1)
+            for i in [1, 2]
+        ]
+        super().deposit(deposit_amounts, user)
+
     @rule(
-        exchange_amount_in=exchange_amount_in,
+        token_amount=token_amount,
         exchange_i=exchange_i,
-        exchange_j=exchange_j,
         user=user,
     )
-    def add_liquidity(self, exchange_amount_in, exchange_i, exchange_j, user):
+    def remove_liquidity_one_coin(self, token_amount, exchange_i, user):
 
-        raise
+        super().remove_liquidity_one_coin(
+            token_amount, exchange_i, user, False
+        )
 
 
 def test_admin_fee(swap, views_contract, users, pool_coins, tricrypto_factory):

tests/boa/unitary/pool/stateful/test_gas_realistic.py

@@ -56,9 +56,12 @@ def deposit(self, deposit_amount, exchange_i, user):
 
             amounts[exchange_i] = deposit_amount
 
-        new_balances = [x + y for x, y in zip(self.balances, amounts)]
         mint_for_testing(self.coins[exchange_i], user, deposit_amount)
 
+        with boa.env.prank(user):
+            for coin in self.coins:
+                coin.approve(self.swap, 2**256 - 1)
+
         try:
 
             tokens = self.token.balanceOf(user)
@@ -68,7 +71,9 @@ def deposit(self, deposit_amount, exchange_i, user):
 
             tokens = self.token.balanceOf(user) - tokens
             self.total_supply += tokens
-            self.balances = new_balances
+
+            for i in range(3):
+                self.balances[i] += amounts[i]
 
         except Exception:
 

tests/boa/unitary/pool/stateful/test_ramp.py

@@ -35,6 +35,8 @@ def setup(self, user_id=0):
         with boa.env.prank(self.tricrypto_factory.admin()):
             self.swap.ramp_A_gamma(new_A, new_gamma, block_time + 14 * 86400)
 
+        self.xcp_profit_a_init = self.swap.xcp_profit_a()
+
     @rule(deposit_amounts=deposit_amounts, user=user)
     def deposit(self, deposit_amounts, user):
         deposit_amounts[1:] = [deposit_amounts[0]] + [
@@ -99,6 +101,17 @@ def virtual_price(self):
         # Invariant is not conserved here
         pass
 
+    @invariant()
+    def check_bumped_fee_during_ramp(self):
+        if self.swap.future_A_gamma_time() > boa.env.vm.state.timestamp:
+            assert (
+                self.swap.fee() >= 10**8
+            )  # Charge at least 100 basis points!
+
+    @invariant()
+    def check_xcp_profit_a_doesnt_increase(self):
+        assert self.swap.xcp_profit_a() == self.xcp_profit_a_init
+
 
 def test_ramp(swap, views_contract, users, pool_coins, tricrypto_factory):
     from hypothesis import settings

tests/boa/unitary/pool/stateful/test_ramp_nocheck.py

@@ -46,6 +46,8 @@ def initialize(self, future_A, future_gamma):
                 boa.env.vm.state.timestamp + 14 * 86400,
             )
 
+        self.xcp_profit_a_init = self.swap.xcp_profit_a()
+
     @rule(
         exchange_amount_in=exchange_amount_in,
         exchange_i=exchange_i,
@@ -75,6 +77,17 @@ def up_only_profit(self):
         # so we need to override super().up_only_profit()
         pass
 
+    @invariant()
+    def check_bumped_fee_during_ramp(self):
+        if self.swap.future_A_gamma_time() > boa.env.vm.state.timestamp:
+            assert (
+                self.swap.fee() >= 10**8
+            )  # Charge at least 100 basis points!
+
+    @invariant()
+    def check_xcp_profit_a_doesnt_increase(self):
+        assert self.swap.xcp_profit_a() == self.xcp_profit_a_init
+
 
 def test_ramp(swap, views_contract, users, pool_coins, tricrypto_factory):
     from hypothesis import settings

tests/boa/unitary/pool/stateful/test_stateful.py

@@ -30,19 +30,23 @@ def deposit(self, deposit_amounts, user):
             return
 
         amounts = self.convert_amounts(deposit_amounts)
-        new_balances = [x + y for x, y in zip(self.balances, amounts)]
 
         for coin, q in zip(self.coins, amounts):
             mint_for_testing(coin, user, q)
 
         try:
+
             tokens = self.token.balanceOf(user)
             with boa.env.prank(user), self.upkeep_on_claim():
                 self.swap.add_liquidity(amounts, 0)
             tokens = self.token.balanceOf(user) - tokens
             self.total_supply += tokens
-            self.balances = new_balances
+
+            for i in range(3):
+                self.balances[i] += amounts[i]
+
         except Exception:
+
             if self.check_limits(amounts):
                 raise
             else: