precise tests

contracts/main/CurveStableSwapMetaNG.vy

@@ -752,6 +752,7 @@ def add_liquidity(
     @param _receiver Address that owns the minted LP tokens
     @return Amount of LP tokens received by depositing
     """
+    amounts: DynArray[uint256, MAX_COINS] = [_amounts[0], _amounts[1]]
     amp: uint256 = self._A()
     old_balances: DynArray[uint256, MAX_COINS] = self._balances()
     rates: DynArray[uint256, MAX_COINS] = self._stored_rates()
@@ -766,12 +767,12 @@ def add_liquidity(
 
     for i in range(N_COINS_128):
 
-        if _amounts[i] > 0:
+        if amounts[i] > 0:
 
             new_balances[i] += self._transfer_in(
                 i,
                 -1,  # <--- we're not handling underlying coins here
-                _amounts[i],
+                amounts[i],
                 msg.sender,
                 False,  # expect_optimistic_transfer
             )
@@ -859,7 +860,7 @@ def add_liquidity(
 
     log AddLiquidity(
         msg.sender,
-        [_amounts[0], _amounts[1]],
+        amounts,
         fees,
         D1,
         total_supply
@@ -924,6 +925,7 @@ def remove_liquidity_imbalance(
     @param _receiver Address that receives the withdrawn coins
     @return Actual amount of the LP token burned in the withdrawal
     """
+    amounts: DynArray[uint256, MAX_COINS] = [_amounts[0], _amounts[1]]
     amp: uint256 = self._A()
     rates: DynArray[uint256, MAX_COINS] = self._stored_rates()
     old_balances: DynArray[uint256, MAX_COINS] = self._balances()
@@ -932,9 +934,9 @@ def remove_liquidity_imbalance(
 
     for i in range(N_COINS_128):
 
-        if _amounts[i] != 0:
-            new_balances[i] -= _amounts[i]
-            self._transfer_out(i, _amounts[i], _receiver)
+        if amounts[i] != 0:
+            new_balances[i] -= amounts[i]
+            self._transfer_out(i, amounts[i], _receiver)
 
     D1: uint256 = self.get_D_mem(rates, new_balances, amp)
     # base_fee: uint256 = self.fee * N_COINS / (4 * (N_COINS - 1))
@@ -984,7 +986,7 @@ def remove_liquidity_imbalance(
 
     log RemoveLiquidityImbalance(
         msg.sender,
-        [_amounts[0], _amounts[1]],
+        amounts,
         fees,
         D1,
         total_supply

tests/pools/meta/test_meta_zap.py

@@ -91,20 +91,21 @@ def zap(base_pool, base_pool_tokens, base_pool_lp_token):
 
 
 @pytest.fixture(scope="module")
-def swap(zap, base_pool, empty_swap, charlie, tokens_all):
+def initial_amts():
+    return [100 * 10**18] * 4
+
+
+@pytest.fixture(scope="module")
+def swap(zap, base_pool, empty_swap, charlie, tokens_all, initial_amts):
 
     for i in range(3):
         assert base_pool.balances(i) == 0
 
-    deposit_amount = 100 * 10**18
-
     for token in tokens_all:
-        mint_for_testing(charlie, deposit_amount, token, False)
+        mint_for_testing(charlie, initial_amts[0], token, False)
         token.approve(zap.address, 2**256 - 1, sender=charlie)
 
-    deposit_amounts = [deposit_amount] * 4
-
-    out_amount = zap.add_liquidity(empty_swap.address, deposit_amounts, 0, sender=charlie)
+    out_amount = zap.add_liquidity(empty_swap.address, initial_amts, 0, sender=charlie)
     assert out_amount > 0
     assert 0 not in empty_swap.get_balances()
     assert empty_swap.totalSupply() > 0
@@ -128,24 +129,26 @@ def test_calc_amts_add(zap, swap, charlie, tokens_all):
     assert calc_amt_zap == out_amount
 
 
-def test_calc_amts_remove_imbalance(zap, swap, charlie, tokens_all):
-
-    to_receive_amounts = [10 * 10**18] * 4
-
-    charlie_bal_before = []
-    for token in tokens_all:
-        charlie_bal_before.append(token.balanceOf(charlie))
+def test_calc_amts_remove_imbalance(
+    zap, swap, meta_token, base_pool_tokens, base_pool_lp_token, base_pool, charlie, tokens_all, initial_amts
+):
 
+    amounts = [i // 4 for i in initial_amts]
+    initial_balance = swap.balanceOf(charlie)
     swap.approve(zap, 2**256 - 1, sender=charlie)
-    calc_burnt_amt_zap = zap.calc_token_amount(swap.address, to_receive_amounts, False)
-    actual_burnt_amt = zap.remove_liquidity_imbalance(
-        swap.address, [int(0.9 * amt) for amt in to_receive_amounts], calc_burnt_amt_zap, sender=charlie
-    )
-
-    assert actual_burnt_amt <= calc_burnt_amt_zap
+    max_burn = swap.balanceOf(charlie)
+    zap.remove_liquidity_imbalance(swap, amounts, max_burn, sender=charlie)
 
+    # check if charlie received what was wanted:
     for i, token in enumerate(tokens_all):
-        assert token.balanceOf(charlie) > charlie_bal_before[i]
+        assert token.balanceOf(charlie) == amounts[i]
+
+    # bob is the only LP, total supply is affected in the same way as his balance
+    assert swap.balanceOf(charlie) < initial_balance
+    assert swap.balanceOf(charlie) >= initial_balance - max_burn
+
+    assert swap.balanceOf(zap) == 0
+    assert swap.balanceOf(charlie) == swap.totalSupply()
 
 
 def test_calc_amts_remove(zap, swap, charlie, tokens_all, meta_token, base_pool, base_pool_tokens):