Skip onchain balance check (#881)

Use same `check_balances` arg in `adjust_spend()` to skip balance check if needed

tests/ethereum/test_execute_trade_uniswap.py

@@ -533,3 +533,87 @@ def test_two_parallel_positions(
     assert balances[asset_aave.address] == 0
     assert balances[asset_weth.address] == 0
 
+
+def test_execute_alpha_model_rebalance_trades(
+    web3: Web3,
+    state: State,
+    uniswap_v2: UniswapV2Deployment,
+    hot_wallet: HotWallet,
+    weth_usdc_pair: TradingPairIdentifier,
+    aave_usdc_pair: TradingPairIdentifier,
+    asset_aave,
+    asset_weth,
+    asset_usdc,
+    pair_universe,
+    start_ts: datetime.datetime,
+    tx_builder,
+):
+    """Execute rebalance trades on 2 positions in parallel:
+    
+    1. Buy 4500 USDC worth of WETH at 1700 USD
+    2. Buy 4500 USDC worth of AAVE at 200 USD
+    3. Sell all WETH (worth 4500 USDC)
+    4. Buy 4500 USDC worth of AAVE
+    """
+
+    portfolio = state.portfolio
+
+    # We have everything in cash and initial assumptions on the price
+    assert portfolio.get_total_equity() == 10_000
+    assert portfolio.get_cash() == 10_000
+    assert get_current_price(web3, uniswap_v2, weth_usdc_pair) == pytest.approx(1693.211867)
+    assert get_current_price(web3, uniswap_v2, aave_usdc_pair) == pytest.approx(199.201396)
+    assert hot_wallet.current_nonce == 0
+
+    #
+    # 1. Buy 4500 USDC worth of WETH at 1700 USD
+    # 2. Buy 4500 USDC worth of AAVE at 200 USD
+    #
+
+    trader = UniswapV2TestTrader(uniswap_v2, state, pair_universe, tx_builder)
+    position1, trade1 = trader.buy(weth_usdc_pair, Decimal(4500), execute=False)
+    position2, trade2 = trader.buy(aave_usdc_pair, Decimal(4500), execute=False)
+
+    assert position1.position_id == 1
+    assert position2.position_id == 2
+
+    assert len(portfolio.open_positions) == 2
+
+    # Execute both trades
+    trader.execute_trades_simple(trader.create_routing_model(), [trade1, trade2])
+    assert hot_wallet.current_nonce == 3
+
+    assert position1.get_quantity_old() == pytest.approx(Decimal(2.632171037438914461))
+    assert position2.get_quantity_old() == pytest.approx(Decimal(21.940323684081220533))
+    assert position1.get_value() == pytest.approx(4500)
+    assert position2.get_value() == pytest.approx(4500)
+    assert portfolio.get_total_equity() == pytest.approx(9999.999998002779)
+    assert portfolio.get_cash() == pytest.approx(1000)
+
+    balances = get_held_assets(web3, hot_wallet.address, [asset_usdc, asset_aave, asset_weth])
+    assert balances[asset_usdc.address] == pytest.approx(Decimal(1000))
+    assert balances[asset_weth.address] == pytest.approx(Decimal(2.632171037438914461))
+    assert balances[asset_aave.address] == pytest.approx(Decimal(21.940323684081220533))
+
+    #
+    # 3. Sell all WETH (worth 4500 USDC)
+    # 4. Use 4500 USDC worth to buy AAVE
+    #
+
+    position3, trade3 = trader.sell(weth_usdc_pair, position1.get_quantity_old(), execute=False)
+    position4, trade4 = trader.buy(aave_usdc_pair, 4500, execute=False)
+
+    trader.execute_trades_simple(trader.create_routing_model(), [trade3, trade4])
+
+    assert trade3.blockchain_transactions[0].nonce == 3
+    assert trade4.blockchain_transactions[0].nonce == 5
+
+    assert position3.is_closed()
+    assert position4.is_open()
+    assert portfolio.get_total_equity() == pytest.approx(9779.609513697993)
+    assert portfolio.get_cash() == pytest.approx(973.11125)
+
+    balances = get_held_assets(web3, hot_wallet.address, [asset_usdc, asset_aave, asset_weth])
+    assert balances[asset_usdc.address] == pytest.approx(Decimal(973.11125))
+    assert balances[asset_weth.address] == 0
+    assert balances[asset_aave.address] == pytest.approx(Decimal(42.937205))

tradeexecutor/ethereum/routing_model.py

@@ -112,6 +112,7 @@ def make_direct_trade(self,
         adjusted_reserve_amount = routing_state.adjust_spend(
             reserve_asset,
             reserve_amount,
+            check_balances=check_balances,
         )
 
         logger.info(
@@ -178,6 +179,7 @@ def make_multihop_trade(self,
         adjusted_reserve_amount = routing_state.adjust_spend(
             reserve_asset,
             reserve_amount,
+            check_balances=check_balances,
         )
 
         trade_txs = routing_state.trade_on_router_three_way(

tradeexecutor/ethereum/routing_state.py

@@ -244,10 +244,11 @@ def create_signed_transaction(
         return [signed_tx]
 
     def adjust_spend(
-            self,
-            asset: AssetIdentifier,
-            required_amount: int,
-            epsilon=0.00001,
+        self,
+        asset: AssetIdentifier,
+        required_amount: int,
+        epsilon: float = 0.00001,
+        check_balances: bool = False,
     ) -> int:
         """Check that our on-chain balances have enough tokens to cover the trade.
 
@@ -274,12 +275,16 @@ def adjust_spend(
         holding_address = self.tx_builder.get_erc_20_balance_address()
         token = fetch_erc20_details(web3, asset.address)
         on_chain_balance = token.contract.functions.balanceOf(holding_address).call()
-        if on_chain_balance <  required_amount:
+        if on_chain_balance < required_amount:
             # Check if we are within epsilon
             if (required_amount - on_chain_balance) / required_amount < epsilon:
                 logger.info("Adjusting spending amount to fit to the epsilon. For %s we have on-chain: %d, required: %d", asset, on_chain_balance, required_amount)
                 return on_chain_balance
             else:
+                if not check_balances:
+                    logger.info("For %s we have on-chain: %d < required: %d, but we skip balance check since check_balances=False", asset, on_chain_balance, required_amount)
+                    return required_amount
+
                 raise OutOfBalance(
                     f"Not enough tokens for {asset} to perform the trade. Required: {required_amount}, on-chain balance for {holding_address} is {on_chain_balance}."
                 )