From 2d5e14d3eb891e655281700985a16b7acddfe270 Mon Sep 17 00:00:00 2001
From: Evan Griffiths <56087052+evangriffiths@users.noreply.github.com>
Date: Tue, 11 Jun 2024 20:54:14 +0100
Subject: [PATCH] Fix autodeposit behaviour inside omen_buy_outcome_tx (#273)

---
 .../markets/omen/omen.py                      | 20 +++++++-
 .../tools/balances.py                         |  4 ++
 tests_integration/markets/omen/test_omen.py   | 51 ++++++++++++++++++-
 3 files changed, 72 insertions(+), 3 deletions(-)

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,
+    )