Use in test

Author: michaelhly

tests/pools/common/unitary/test_remove_liquidity_one_coin.py

@@ -1,5 +1,6 @@
 import brownie
 import pytest
+from simulation import Curve
 
 pytestmark = [
     pytest.mark.skip_pool("busd", "compound", "pax", "susd", "usdt", "y"),
@@ -29,20 +30,23 @@ def test_amount_received(chain, alice, swap, wrapped_coins, wrapped_decimals, id
     else:
         assert 10**decimals // rate_mod <= wrapped.balanceOf(alice) <= 10**decimals
 
-
 @pytest.mark.itercoins("idx")
 @pytest.mark.parametrize("divisor", [1, 5, 42])
-def test_lp_token_balance(alice, swap, pool_token, idx, divisor, n_coins, base_amount):
+def test_lp_token_balance(alice, swap, wrapped_coins, pool_token, idx, divisor, n_coins, base_amount):
     amount = pool_token.balanceOf(alice) // divisor
 
     swap.remove_liquidity_one_coin(amount, idx, 0, {'from': alice})
 
     assert pool_token.balanceOf(alice) == n_coins * 10**18 * base_amount - amount
 
+    balances = [swap.balances(i) for i in range(n_coins)]
+    curve_model = Curve(swap.A(), balances, n_coins, tokens=pool_token.totalSupply())
+    assert wrapped_coins[idx].balanceOf(alice) == curve_model.calc_withdraw_one_coin(amount, idx)[0]
+
 
 @pytest.mark.itercoins("idx")
 @pytest.mark.parametrize("rate_mod", [0.9, 1.1])
-def test_expected_vs_actual(chain, alice, swap, wrapped_coins, pool_token, idx, rate_mod):
+def test_expected_vs_actual(chain, alice, swap, wrapped_coins, pool_token, idx, n_coins, rate_mod):
     amount = pool_token.balanceOf(alice) // 10
     wrapped = wrapped_coins[idx]
 
@@ -57,6 +61,10 @@ def test_expected_vs_actual(chain, alice, swap, wrapped_coins, pool_token, idx,
 
     assert wrapped_coins[idx].balanceOf(alice) == expected
 
+    balances = [swap.balances(i) for i in range(n_coins)]
+    curve_model = Curve(swap.A(), balances, n_coins, tokens=pool_token.totalSupply())
+    assert wrapped_coins[idx].balanceOf(alice) == curve_model.calc_withdraw_one_coin(amount, idx)[0]
+
 
 @pytest.mark.itercoins("idx")
 def test_below_min_amount(alice, swap, wrapped_coins, pool_token, idx):