diff --git a/prediction_market_agent_tooling/markets/omen/omen.py b/prediction_market_agent_tooling/markets/omen/omen.py index d88ade90..352855fd 100644 --- a/prediction_market_agent_tooling/markets/omen/omen.py +++ b/prediction_market_agent_tooling/markets/omen/omen.py @@ -142,6 +142,7 @@ def place_bet( amount: BetAmount, omen_auto_deposit: bool = True, web3: Web3 | None = None, + api_keys: APIKeys | None = None, ) -> None: if not self.can_be_traded(): raise ValueError( @@ -151,7 +152,7 @@ def place_bet( raise ValueError(f"Omen bets are made in xDai. Got {amount.currency}.") amount_xdai = xDai(amount.amount) binary_omen_buy_outcome_tx( - api_keys=APIKeys(), + api_keys=api_keys if api_keys is not None else APIKeys(), amount=amount_xdai, market=self, binary_outcome=outcome, @@ -159,6 +160,20 @@ def place_bet( web3=web3, ) + def buy_tokens( + self, + outcome: bool, + amount: TokenAmount, + web3: Web3 | None = None, + api_keys: APIKeys | None = None, + ) -> None: + return self.place_bet( + outcome=outcome, + amount=amount, + web3=web3, + api_keys=api_keys, + ) + def calculate_sell_amount_in_collateral( self, amount: TokenAmount, outcome: bool, web3: Web3 | None = None ) -> xDai: @@ -460,8 +475,9 @@ def omen_buy_outcome_tx( for_address=from_address_checksummed, web3=web3 ) if auto_deposit and collateral_token_balance < amount_wei: + deposit_amount_wei = Wei(amount_wei - collateral_token_balance) collateral_token_contract.deposit( - api_keys=api_keys, amount_wei=amount_wei, web3=web3 + api_keys=api_keys, amount_wei=deposit_amount_wei, web3=web3 ) # Buy shares using the deposited xDai in the collateral token. market_contract.buy( diff --git a/prediction_market_agent_tooling/tools/balances.py b/prediction_market_agent_tooling/tools/balances.py index 05bcc4c1..6622b3ed 100644 --- a/prediction_market_agent_tooling/tools/balances.py +++ b/prediction_market_agent_tooling/tools/balances.py @@ -13,6 +13,10 @@ class Balances(BaseModel): xdai: xDai wxdai: xDai + @property + def total(self) -> xDai: + return xDai(self.xdai + self.wxdai) + def get_balances(address: ChecksumAddress, web3: Web3 | None = None) -> Balances: if not web3: diff --git a/tests_integration/markets/omen/test_omen.py b/tests_integration/markets/omen/test_omen.py index b647f43e..f7b8e9ba 100644 --- a/tests_integration/markets/omen/test_omen.py +++ b/tests_integration/markets/omen/test_omen.py @@ -10,7 +10,11 @@ from prediction_market_agent_tooling.config import APIKeys from prediction_market_agent_tooling.gtypes import xDai, xdai_type from prediction_market_agent_tooling.loggers import logger -from prediction_market_agent_tooling.markets.data_models import Currency, TokenAmount +from prediction_market_agent_tooling.markets.data_models import ( + BetAmount, + Currency, + TokenAmount, +) from prediction_market_agent_tooling.markets.omen.data_models import ( OMEN_FALSE_OUTCOME, OMEN_TRUE_OUTCOME, @@ -26,6 +30,7 @@ pick_binary_market, ) from prediction_market_agent_tooling.markets.omen.omen_contracts import ( + OmenCollateralTokenContract, OmenRealitioContract, ) from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import ( @@ -259,3 +264,47 @@ def get_market_outcome_tokens() -> TokenAmount: # Check that we have sold our entire stake in the market. remaining_tokens = get_market_outcome_tokens() assert np.isclose(remaining_tokens.amount, 0, atol=1e-6) + + +def test_place_bet_with_autodeposit( + local_web3: Web3, + test_keys: APIKeys, +) -> None: + market = OmenAgentMarket.from_data_model(pick_binary_market()) + initial_balances = get_balances(address=test_keys.bet_from_address, web3=local_web3) + collateral_token_contract = OmenCollateralTokenContract() + + # Start by moving all funds from wxdai to xdai + if initial_balances.wxdai > 0: + collateral_token_contract.withdraw( + api_keys=test_keys, + amount_wei=xdai_to_wei(initial_balances.wxdai), + web3=local_web3, + ) + + # Check that we have xdai funds, but no wxdai funds + initial_balances = get_balances(address=test_keys.bet_from_address, web3=local_web3) + assert initial_balances.wxdai == 0 + assert initial_balances.xdai > 0 + + # Convert half of the xDai to wxDai + collateral_token_contract.deposit( + api_keys=test_keys, + amount_wei=xdai_to_wei(initial_balances.xdai * 0.5), + web3=local_web3, + ) + new_balances = get_balances(address=test_keys.bet_from_address, web3=local_web3) + assert np.allclose(new_balances.total, initial_balances.total) + + # Try to place a bet with 90% of the xDai funds + bet_amount = BetAmount(amount=initial_balances.xdai * 0.9, currency=Currency.xDai) + assert new_balances.xdai < bet_amount.amount + assert new_balances.wxdai < bet_amount.amount + assert new_balances.total > bet_amount.amount + market.place_bet( + outcome=True, + amount=bet_amount, + omen_auto_deposit=True, + web3=local_web3, + api_keys=test_keys, + )