Added bet_from_address for Safe support | added tests

prediction_market_agent_tooling/markets/omen/omen.py

@@ -42,14 +42,14 @@
     PRESAGIO_BASE_URL,
     Condition,
     ConditionPreparationEvent,
-    ContractPrediction,
     CreatedMarket,
     IPFSAgentResult,
     OmenBet,
     OmenMarket,
     OmenUserPosition,
     get_bet_outcome,
     get_boolean_outcome,
+    ContractPrediction,
 )
 from prediction_market_agent_tooling.markets.omen.omen_contracts import (
     OMEN_DEFAULT_MARKET_FEE_PERC,
@@ -79,10 +79,10 @@
 from prediction_market_agent_tooling.tools.hexbytes_custom import HexBytes
 from prediction_market_agent_tooling.tools.ipfs.ipfs_handler import IPFSHandler
 from prediction_market_agent_tooling.tools.utils import (
-    BPS_CONSTANT,
     DatetimeUTC,
     calculate_sell_amount_in_collateral,
     check_not_none,
+    BPS_CONSTANT,
 )
 from prediction_market_agent_tooling.tools.web3_utils import (
     add_fraction,
@@ -451,7 +451,7 @@ def store_trades(
             HexBytes(HexStr(i.id)) for i in traded_market.trades if i.id is not None
         ]
         prediction = ContractPrediction(
-            publisher=keys.public_key,
+            publisher=keys.bet_from_address,
             ipfs_hash=ipfs_hash_decoded,
             tx_hashes=tx_hashes,
             estimated_probability_bps=int(traded_market.answer.p_yes * BPS_CONSTANT),

tests_integration_with_local_chain/markets/omen/test_omen.py

@@ -462,7 +462,7 @@ def get_position_balance_by_position_id(
 )
 def test_add_predictions(local_web3: Web3, test_keys: APIKeys, ipfs_hash: str) -> None:
     agent_result_mapping = OmenAgentResultMappingContract()
-    market_address = test_keys.public_key
+    market_address = test_keys.bet_from_address
     dummy_transaction_hash = "0x3750ffa211dab39b4d0711eb27b02b56a17fa9d257ee549baa3110725fd1d41b"  # web3-private-key-ok
     stored_predictions = agent_result_mapping.get_predictions(
         market_address, web3=local_web3
@@ -471,7 +471,7 @@ def test_add_predictions(local_web3: Web3, test_keys: APIKeys, ipfs_hash: str) -
         tx_hashes=[HexBytes(dummy_transaction_hash)],
         estimated_probability_bps=5454,
         ipfs_hash=HexBytes(ipfs_hash),
-        publisher=test_keys.public_key,
+        publisher=test_keys.bet_from_address,
     )
 
     agent_result_mapping.add_prediction(test_keys, market_address, p, web3=local_web3)

tests_integration_with_local_chain/safe/test_safe.py

@@ -3,21 +3,29 @@
 from safe_eth.eth import EthereumClient
 from safe_eth.safe.safe import SafeV141
 from web3 import Web3
+from web3.constants import HASH_ZERO, ADDRESS_ZERO
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.gtypes import PrivateKey, xDai
 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.omen.data_models import OMEN_TRUE_OUTCOME
+from prediction_market_agent_tooling.markets.omen.data_models import (
+    OMEN_TRUE_OUTCOME,
+    ContractPrediction,
+)
 from prediction_market_agent_tooling.markets.omen.omen import (
     FilterBy,
     OmenAgentMarket,
     SortBy,
     binary_omen_buy_outcome_tx,
 )
+from prediction_market_agent_tooling.markets.omen.omen_contracts import (
+    OmenAgentResultMappingContract,
+)
 from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import (
     OmenSubgraphHandler,
 )
+from prediction_market_agent_tooling.tools.hexbytes_custom import HexBytes
 from prediction_market_agent_tooling.tools.web3_utils import (
     send_xdai_to,
     xdai_to_wei,
@@ -98,3 +106,42 @@ def test_send_function_on_contract_tx_using_safe(
     logger.debug(f"final Yes token balance {final_yes_token_balance}")
     print_current_block(local_web3)
     assert initial_yes_token_balance.amount < final_yes_token_balance.amount
+
+
+def test_add_prediction_with_safe(
+    local_ethereum_client: EthereumClient,
+    local_web3: Web3,
+    test_keys: APIKeys,
+    test_safe: SafeV141,
+) -> None:
+    test_keys.SAFE_ADDRESS = test_safe.address
+    dummy_transaction_hash = "0x3750ffa211dab39b4d0711eb27b02b56a17fa9d257ee549baa3110725fd1d41b"  # web3-private-key-ok
+    dummy_market_address = Web3.to_checksum_address(ADDRESS_ZERO)
+    p = ContractPrediction(
+        tx_hashes=[HexBytes(dummy_transaction_hash)],
+        estimated_probability_bps=5454,
+        ipfs_hash=HexBytes(HASH_ZERO),
+        publisher=test_keys.bet_from_address,
+    )
+
+    contract = OmenAgentResultMappingContract()
+
+    tx_receipt = contract.add_prediction(
+        api_keys=test_keys,
+        market_address=dummy_market_address,
+        prediction=p,
+        web3=local_web3,
+    )
+    local_web3.eth.wait_for_transaction_receipt(
+        transaction_hash=tx_receipt.transactionHash
+    )
+    # We expect a new prediction to exist under the Safe address key.
+    predictions = contract.get_predictions(
+        market_address=dummy_market_address, web3=local_web3
+    )
+    predictions_from_safe_address = [
+        i for i in predictions if i.publisher_checksummed == test_keys.bet_from_address
+    ]
+    assert len(predictions_from_safe_address) == 1
+    actual_prediction = predictions_from_safe_address[0]
+    assert actual_prediction == p