Allowing LRNA price in cash_out_omnipool. Fixed test_cash_out_multipl…

…e_positions
galacticcouncil · Sep 20, 2024 · c9b542f · c9b542f
1 parent 5b7838a
commit c9b542f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 11 deletions.
diff --git a/hydradx/model/amm/omnipool_amm.py b/hydradx/model/amm/omnipool_amm.py
@@ -1590,15 +1590,16 @@ def cash_out_omnipool(omnipool: OmnipoolState, agent: Agent, prices) -> float:
     lrna_removed = {tkn: -delta_q[tkn] if tkn in delta_q else 0 for tkn in omnipool.asset_list}
     liquidity_removed = {tkn: -delta_r[tkn] if tkn in delta_r else 0 for tkn in omnipool.asset_list}
 
-    if 'LRNA' in prices:
-        raise ValueError('LRNA price should not be given.')
-    lrna_to_sell = delta_qa
+    # if 'LRNA' in prices:
+    #     raise ValueError('LRNA price should not be given.')
+    agent_lrna = delta_qa
     if 'LRNA' in agent.holdings:
-        lrna_to_sell += agent.holdings['LRNA']
-    if 'LRNA' not in prices and lrna_to_sell > 0:
+        agent_lrna += agent.holdings['LRNA']
+
+    if 'LRNA' not in prices and agent_lrna > 0:
         lrna_total = omnipool.lrna_total - sum(lrna_removed.values())
         lrna_sells = {
-            tkn: -(omnipool.lrna[tkn] - lrna_removed[tkn]) / lrna_total * lrna_to_sell
+            tkn: -(omnipool.lrna[tkn] - lrna_removed[tkn]) / lrna_total * agent_lrna
             for tkn in omnipool.asset_list
         }
 
@@ -1619,6 +1620,7 @@ def cash_out_omnipool(omnipool: OmnipoolState, agent: Agent, prices) -> float:
             liquidity_removed[tkn] += lrna_profits[tkn]
 
     new_holdings = {tkn: agent.holdings[tkn] for tkn in agent.holdings}
+    new_holdings['LRNA'] = agent_lrna
     for tkn in liquidity_removed:
         if tkn not in new_holdings:
             new_holdings[tkn] = 0

diff --git a/hydradx/tests/test_omnipool_amm.py b/hydradx/tests/test_omnipool_amm.py
@@ -3127,17 +3127,19 @@ def test_cash_out_multiple_positions(trade_sizes: list[float]):
         withdrawal_fee=False
     )
 
-    lp_quantity = 10000
+    lp_quantity = mpf(10000)
     agent1 = Agent(holdings={'DOT': lp_quantity * len(trade_sizes)})
-    agent2 = Agent(holdings={'DOT': 10000000, 'HDX': 10000000})
+    agent2 = Agent(holdings={'DOT': mpf(10000000), 'HDX': mpf(10000000)})
     for i, trade in enumerate(trade_sizes):
         initial_state.add_liquidity(agent1, tkn_add='DOT', quantity=lp_quantity, nft_id=str(i))
         if trade > 0:
-            initial_state.swap(agent2, tkn_buy='HDX', tkn_sell='DOT', sell_quantity=trade)
-        else:
-            initial_state.swap(agent2, tkn_buy='DOT', tkn_sell='HDX', sell_quantity=-trade)
+            initial_state.swap(agent2, tkn_buy='HDX', tkn_sell='DOT', sell_quantity=mpf(trade))
+        elif trade < 0:
+            initial_state.swap(agent2, tkn_buy='DOT', tkn_sell='HDX', sell_quantity=-mpf(trade))
 
     spot_prices = {tkn: initial_state.price(initial_state, tkn, 'USD') for tkn in initial_state.asset_list}
+    spot_prices['LRNA'] = oamm.usd_price(initial_state, 'LRNA')
+
     cash_out_value = cash_out_omnipool(initial_state, agent1, spot_prices)
     cash_out_state = initial_state.copy()
     cash_out_agent = agent1.copy()