Integrated pytest

trader/arbitrage.py

@@ -1,31 +1,76 @@
-from typing import List
+"""Module for identifying arbitrage opportunities among pool data."""
+
+from typing import List, Optional, Tuple
+
 from pool_db_data import PoolDBData
-
-def find_arbitrage_opportunities(pool_data: List[PoolDBData]):
+
+
+def find_arbitrage_opportunities(
+    pool_data: List[PoolDBData]
+) -> Optional[Tuple[str, float, str, float, float]]:
+    """
+    Identify arbitrage opportunities based on pool prices.
+
+    This function filters valid pools, extracts their prices, identifies the pools
+    with the minimum and maximum prices, and determines if a profitable arbitrage
+    opportunity exists between them.
+
+    Args:
+        pool_data (List[PoolDBData]): A list of PoolDBData instances representing
+                                      different liquidity pools.
+
+    Returns:
+        Optional[Tuple[str, float, str, float, float]]:
+            - A tuple containing:
+                1. RPC address to buy from (str)
+                2. Buy price (float)
+                3. RPC address to sell to (str)
+                4. Sell price (float)
+                5. Potential margin per base token (float)
+            - Returns None if no arbitrage opportunity is found.
+    """
     # Filter only valid pools
-    valid_pools = [p for p in pool_data if p.isValid]
+    valid_pools = [pool for pool in pool_data if pool.is_valid]
+    print(f"Valid pools count: {len(valid_pools)}")
 
     # Extract prices as floats
-    # Note: poolPrice is assumed to be the direct ratio price (quote per base).
-    pool_prices = [(p.rpcData, float(p.poolPrice)) for p in valid_pools]
+    # Note: pool_price is assumed to be the direct ratio price (quote per base).
+    pool_prices = []
+    for pool in valid_pools:
+        price = float(pool.pool_price)
+        pool_prices.append((pool.rpc_data, price))
+        print(f"Processed pool {pool.rpc_data} with price {price}")
 
     if not pool_prices:
-        print("No valid pools available.")
-        return
+        print("No valid pools with valid prices available.")
+        return None
 
-    # Identify min and max price pools
+    # Identify pools with minimum and maximum prices
     min_price_pool = min(pool_prices, key=lambda x: x[1])
     max_price_pool = max(pool_prices, key=lambda x: x[1])
 
     min_price = min_price_pool[1]
     max_price = max_price_pool[1]
 
+    print(f"Minimum price pool: {min_price_pool[0]} at price {min_price}")
+    print(f"Maximum price pool: {max_price_pool[0]} at price {max_price}")
+
     # Check if there's a profitable price difference
     if max_price > min_price:
         # Potential arbitrage margin per base unit
         margin = max_price - min_price
-        print(f"Arbitrage opportunity found:")
-        print(f"Buy from {min_price_pool[0]} at price {min_price} and sell to {max_price_pool[0]} at price {max_price}.")
+        print("Arbitrage opportunity found:")
+        print(
+            f"Buy from {min_price_pool[0]} at price {min_price} and sell to "
+            f"{max_price_pool[0]} at price {max_price}."
+        )
         print(f"Potential margin per base token: {margin}")
-    else:
-        print("No arbitrage opportunity found.")
+        return (
+            min_price_pool[0],
+            min_price,
+            max_price_pool[0],
+            max_price,
+            margin
+        )
+    print("No arbitrage opportunity found.")
+    return None

trader/main.py

@@ -1,7 +1,10 @@
-from arbitrage_test import run_test_scenarios
+"""Rebalance script module."""
 
 def main():
-  run_test_scenarios()
+    """Prints the 'rebalance' message."""
+    print("rebalance")
+
 
 if __name__ == "__main__":
-  main()
+    main()
+

trader/pool_db_data.py

@@ -1,47 +1,81 @@
+"""Module for managing pool database data and public keys."""
+
 from dataclasses import dataclass
-from typing import Union
+
 
 class PublicKey:
     """
     A placeholder class to represent a PublicKey.
+
     Replace this with the actual PublicKey implementation as required.
     """
-    def __init__(self, value: str):
+
+    def __init__(self, value: str) -> None:
+        """
+        Initialize a PublicKey instance.
+
+        Args:
+            value (str): The value of the public key.
+        """
         self.value = value
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Return the official string representation of the PublicKey."""
         return f"PublicKey({self.value})"
 
-    def __str__(self):
+    def __str__(self) -> str:
+        """Return the informal string representation of the PublicKey."""
         return self.value
 
+
+# pylint: disable=too-many-instance-attributes
 @dataclass
 class PoolDBData:
-    rpcData: str
-    baseReserve: str
-    quoteReserve: str
-    mintAAmount: str
-    mintBAmount: str
-    poolPrice: str
-    lastUpdated: int  # Unix timestamp (seconds since epoch)
-    isValid: bool
-    accountId: PublicKey
-    programId: PublicKey
-
-    def __post_init__(self):
+    """
+    Data class representing pool database information.
+
+    Attributes:
+        rpc_data (str): RPC data related to the pool.
+        base_reserve (str): Base reserve amount.
+        quote_reserve (str): Quote reserve amount.
+        mint_a_amount (str): Amount of mint A.
+        mint_b_amount (str): Amount of mint B.
+        pool_price (str): Current pool price.
+        last_updated (int): Unix timestamp of the last update.
+        is_valid (bool): Validity status of the pool data.
+        account_id (PublicKey): Account ID associated with the pool.
+        program_id (PublicKey): Program ID associated with the pool.
+    """
+
+    rpc_data: str
+    base_reserve: str
+    quote_reserve: str
+    mint_a_amount: str
+    mint_b_amount: str
+    pool_price: str
+    last_updated: int  # Unix timestamp (seconds since epoch)
+    is_valid: bool
+    account_id: PublicKey
+    program_id: PublicKey
+
+    def __post_init__(self) -> None:
         """
-        Post-initialization to validate or process any fields if necessary.
+        Post-initialization to validate fields.
+
+        Raises:
+            ValueError: If account_id or program_id is not an instance of PublicKey.
         """
-        if not isinstance(self.accountId, PublicKey):
-            raise ValueError("accountId must be an instance of PublicKey")
-        if not isinstance(self.programId, PublicKey):
-            raise ValueError("programId must be an instance of PublicKey")
+        if not isinstance(self.account_id, PublicKey):
+            raise ValueError("account_id must be an instance of PublicKey")
+        if not isinstance(self.program_id, PublicKey):
+            raise ValueError("program_id must be an instance of PublicKey")
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Return the official string representation of the PoolDBData."""
         return (
-            f"PoolDBData(rpcData={self.rpcData}, baseReserve={self.baseReserve}, "
-            f"quoteReserve={self.quoteReserve}, mintAAmount={self.mintAAmount}, "
-            f"mintBAmount={self.mintBAmount}, poolPrice={self.poolPrice}, "
-            f"lastUpdated={self.lastUpdated}, isValid={self.isValid}, "
-            f"accountId={self.accountId}, programId={self.programId})"
+            f"PoolDBData(rpc_data={self.rpc_data}, base_reserve={self.base_reserve}, "
+            f"quote_reserve={self.quote_reserve}, mint_a_amount={self.mint_a_amount}, "
+            f"mint_b_amount={self.mint_b_amount}, pool_price={self.pool_price}, "
+            f"last_updated={self.last_updated}, is_valid={self.is_valid}, "
+            f"account_id={self.account_id}, program_id={self.program_id})"
         )

trader/requirements.txt

@@ -1,2 +1,3 @@
 networkx==3.4.1
 redis==5.0.0
+pytest==8.3.4