Partial restructuring of repo

src/daemon.py

@@ -16,10 +16,6 @@ def main() -> None:
     process_fees = False
     process_prices = True
 
-    web3, db_engine = initialize_connections()
-    blockchain = BlockchainData(web3)
-    db = Database(db_engine, chain_name)
-
     if chain_name == "arbitrum_one":
         process_imbalances = False
         process_prices = False
@@ -28,11 +24,10 @@ def main() -> None:
         process_prices = False
 
     processor = TransactionProcessor(
-        blockchain, db, chain_name, process_imbalances, process_fees, process_prices
+        chain_name, process_imbalances, process_fees, process_prices
     )
 
-    start_block = processor.get_start_block()
-    processor.process(start_block)
+    processor.run()
 
 
 if __name__ == "__main__":

src/imbalances_script.py → src/raw_imbalances.py

@@ -297,7 +297,7 @@ def filter_sdai_events(event_list: list[dict], is_deposit: bool) -> None:
         filter_sdai_events(events["DepositSDAI"], is_deposit=True)
         filter_sdai_events(events["WithdrawSDAI"], is_deposit=False)
 
-    def compute_imbalances(self, tx_hash: str) -> dict[str, int]:
+    def aggregate_imbalances(self, tx_hash: str) -> dict[str, int]:
         try:
             tx_receipt = self.get_transaction_receipt(tx_hash)
             if not tx_receipt:
@@ -332,14 +332,32 @@ def compute_imbalances(self, tx_hash: str) -> dict[str, int]:
             logger.error("Error computing imbalances for %s: %s", tx_hash, e)
             raise
 
+    def compute_token_imbalances(self, tx_hash: str) -> dict[str, int]:
+        """Process token imbalances for a given transaction and return imbalances."""
+        try:
+            token_imbalances = self.aggregate_imbalances(tx_hash)
+            if token_imbalances:
+                logger.info(f"Token Imbalances on {self.chain_name} for tx {tx_hash}:")
+            return token_imbalances
+        except Exception as e:
+            logger.error(f"Failed to compute imbalances for transaction {tx_hash}: {e}")
+            return {}
+
+    def get_transaction_tokens(self, tx_hash: str) -> list[tuple[str, str]]:
+        token_imbalances = self.compute_token_imbalances(tx_hash)
+        transaction_tokens = []
+        for token_address in token_imbalances.keys():
+            transaction_tokens.append((tx_hash, token_address))
+        return transaction_tokens
+
 
 def main() -> None:
     """main function for finding imbalance for a single tx hash."""
     tx_hash = input("Enter transaction hash: ")
     chain_name, web3 = find_chain_with_tx(tx_hash)
     rt = RawTokenImbalances(web3, chain_name)
     try:
-        imbalances = rt.compute_imbalances(tx_hash)
+        imbalances = rt.aggregate_imbalances(tx_hash)
         if imbalances:
             logger.info(f"Token Imbalances on {chain_name}:")
             for token_address, imbalance in imbalances.items():

src/test_single_hash.py

@@ -1,6 +1,6 @@
 from hexbytes import HexBytes
 from web3 import Web3
-from src.imbalances_script import RawTokenImbalances
+from src.raw_imbalances import RawTokenImbalances
 from src.price_providers.price_feed import PriceFeed
 from src.fees.compute_fees import compute_all_fees_of_batch
 from src.transaction_processor import calculate_slippage
@@ -25,7 +25,7 @@ def __init__(self):
         self.price_providers = PriceFeed()
 
     def compute_data(self, tx_hash: str):
-        token_imbalances = self.imbalances.compute_imbalances(tx_hash)
+        token_imbalances = self.imbalances.compute_token_imbalances(tx_hash)
         protocol_fees, partner_fees, network_fees = compute_all_fees_of_batch(
             HexBytes(tx_hash)
         )

src/transaction_processor.py

@@ -5,10 +5,10 @@
 
 from src.fees.compute_fees import compute_all_fees_of_batch
 from src.helpers.blockchain_data import BlockchainData
-from src.helpers.config import CHAIN_SLEEP_TIME, logger
+from src.helpers.config import CHAIN_SLEEP_TIME, logger, initialize_connections
 from src.helpers.database import Database
-from src.helpers.helper_functions import read_sql_file, set_params
-from src.imbalances_script import RawTokenImbalances
+from src.helpers.helper_functions import set_params
+from src.raw_imbalances import RawTokenImbalances
 from src.price_providers.price_feed import PriceFeed
 from src.token_decimals import update_token_decimals
 
@@ -22,20 +22,19 @@ class TransactionProcessor:
 
     def __init__(
         self,
-        blockchain_data: BlockchainData,
-        db: Database,
         chain_name: str,
         process_imbalances: bool,
         process_fees: bool,
         process_prices: bool,
     ):
-        self.blockchain_data = blockchain_data
-        self.db = db
         self.chain_name = chain_name
         self.process_imbalances = process_imbalances
         self.process_fees = process_fees
         self.process_prices = process_prices
 
+        web3, db_engine = initialize_connections()
+        self.blockchain_data = BlockchainData(web3)
+        self.db = Database(db_engine, chain_name)
         self.imbalances = RawTokenImbalances(self.blockchain_data.web3, self.chain_name)
         self.price_providers = PriceFeed(activate=process_prices)
         self.log_message: list[str] = []
@@ -81,8 +80,11 @@ def get_start_block(self) -> int:
 
         return start_block
 
-    def process(self, start_block: int) -> None:
-        """Main Daemon loop that finds imbalances for txs and prices."""
+    def run(self) -> None:
+        """Main Daemon loop that processes txs and computes imbalances,
+        relevant prices and fees, if needed."""
+
+        start_block = self.get_start_block()
         previous_block = start_block
         unprocessed_txs: list[tuple[str, int, int]] = []
         logger.info("%s daemon started. Start block: %d", self.chain_name, start_block)
@@ -118,11 +120,6 @@ def process_single_transaction(
         """Function processes a single tx to find imbalances, fees, prices including writing to database."""
         self.log_message = []
         try:
-            # compute raw token imbalances
-            token_imbalances = self.process_token_imbalances(
-                tx_hash, auction_id, block_number
-            )
-
             # get transaction timestamp
             transaction_timestamp = self.blockchain_data.get_transaction_timestamp(
                 tx_hash
@@ -131,12 +128,8 @@ def process_single_transaction(
             self.db.write_transaction_timestamp(transaction_timestamp)
 
             # get transaction tokens
-            # transaction_tokens = self.blockchain_data.get_transaction_tokens(tx_hash)
+            transaction_tokens = self.imbalances.get_transaction_tokens(tx_hash)
             # store transaction tokens
-            transaction_tokens = []
-            for token_address, imbalance in token_imbalances.items():
-                if imbalance != 0:
-                    transaction_tokens.append((tx_hash, token_address))
             self.db.write_transaction_tokens(transaction_tokens)
 
             # update token decimals
@@ -177,10 +170,12 @@ def process_single_transaction(
             # ):
             #     return
 
-            if self.process_imbalances and token_imbalances:
-                self.handle_imbalances(
-                    token_imbalances, tx_hash, auction_id, block_number
-                )
+            if self.process_imbalances:
+                token_imbalances = self.imbalances.compute_token_imbalances(tx_hash)
+                if token_imbalances:
+                    self.handle_imbalances(
+                        token_imbalances, tx_hash, auction_id, block_number
+                    )
 
             # if self.process_fees:
             #     self.handle_fees(
@@ -201,21 +196,6 @@ def process_single_transaction(
             logger.error(f"An Error occurred: {err}")
             return
 
-    def process_token_imbalances(
-        self, tx_hash: str, auction_id: int, block_number: int
-    ) -> dict[str, int]:
-        """Process token imbalances for a given transaction and return imbalances."""
-        try:
-            token_imbalances = self.imbalances.compute_imbalances(tx_hash)
-            if token_imbalances:
-                self.log_message.append(
-                    f"Token Imbalances on {self.chain_name} for tx {tx_hash}:"
-                )
-            return token_imbalances
-        except Exception as e:
-            logger.error(f"Failed to compute imbalances for transaction {tx_hash}: {e}")
-            return {}
-
     def process_fees_for_transaction(
         self,
         tx_hash: str,

tests/e2e/test_blockchain_data.py

@@ -3,6 +3,7 @@
 from web3 import Web3
 
 from src.helpers.blockchain_data import BlockchainData
+from src.raw_imbalances import RawTokenImbalances
 
 
 def tests_get_tx_hashes_blocks():
@@ -25,22 +26,3 @@ def test_get_transaction_timestamp():
     transaction_timestamp = blockchain.get_transaction_timestamp(tx_hash)
 
     assert transaction_timestamp == (tx_hash, 1728044411)
-
-
-def test_get_transaction_tokens():
-    web3 = Web3(Web3.HTTPProvider(getenv("NODE_URL")))
-    blockchain = BlockchainData(web3)
-    tx_hash = "0xb75e03b63d4f06c56549effd503e1e37f3ccfc3c00e6985a5aacc9b0534d7c5c"
-
-    transaction_tokens = blockchain.get_transaction_tokens(tx_hash)
-
-    assert all(h == tx_hash for h, _ in transaction_tokens)
-    assert set(token_address for _, token_address in transaction_tokens) == {
-        "0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9",
-        "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48",
-        "0x812Ba41e071C7b7fA4EBcFB62dF5F45f6fA853Ee",
-        "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
-        "0xdAC17F958D2ee523a2206206994597C13D831ec7",
-        "0xEE2a03Aa6Dacf51C18679C516ad5283d8E7C2637",
-        "0x72e4f9F808C49A2a61dE9C5896298920Dc4EEEa9",
-    }

tests/e2e/test_imbalances.py

@@ -0,0 +1,40 @@
+from os import getenv
+
+from web3 import Web3
+
+from src.raw_imbalances import RawTokenImbalances
+
+
+def tests_process_single_transaction():
+    web3 = Web3(Web3.HTTPProvider(getenv("NODE_URL")))
+    raw_imbalances = RawTokenImbalances(web3, "mainnet")
+    imbalances = raw_imbalances.compute_token_imbalances(
+        "0xb75e03b63d4f06c56549effd503e1e37f3ccfc3c00e6985a5aacc9b0534d7c5c"
+    )
+
+    assert imbalances == {
+        "0x72e4f9F808C49A2a61dE9C5896298920Dc4EEEa9": 3116463005,
+        "0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9": 31552225710415395,
+        "0x812Ba41e071C7b7fA4EBcFB62dF5F45f6fA853Ee": 0,
+        "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48": 0,
+        "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2": 0,
+        "0xEE2a03Aa6Dacf51C18679C516ad5283d8E7C2637": 275548164523,
+        "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE": 0,
+    }
+
+def test_get_transaction_tokens():
+    web3 = Web3(Web3.HTTPProvider(getenv("NODE_URL")))
+    imbalances = RawTokenImbalances(web3, "mainnet")
+    tx_hash = "0xb75e03b63d4f06c56549effd503e1e37f3ccfc3c00e6985a5aacc9b0534d7c5c"
+
+    transaction_tokens = imbalances.get_transaction_tokens(tx_hash)
+
+    assert all(h == tx_hash for h, _ in transaction_tokens)
+    assert set(token_address for _, token_address in transaction_tokens) == {
+        "0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9",
+        "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48",
+        "0x812Ba41e071C7b7fA4EBcFB62dF5F45f6fA853Ee",
+        "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
+        "0xEE2a03Aa6Dacf51C18679C516ad5283d8E7C2637",
+        "0x72e4f9F808C49A2a61dE9C5896298920Dc4EEEa9",
+    }

tests/e2e/test_transaction_processor.py

@@ -9,16 +9,12 @@
 
 def tests_process_single_transaction():
     chain_name = "mainnet"
-    web3, db_engine = initialize_connections()
-    blockchain = BlockchainData(web3)
-    db = Database(db_engine, chain_name)
-
     process_imbalances = True
     process_fees = False
     process_prices = True
 
     processor = TransactionProcessor(
-        blockchain, db, chain_name, process_imbalances, process_fees, process_prices
+        chain_name, process_imbalances, process_fees, process_prices
     )
 
     # delete data

tests/legacy/basic_test.py

@@ -2,7 +2,7 @@
 import os
 from dotenv import load_dotenv
 import pytest
-from src.imbalances_script import RawTokenImbalances
+from src.raw_imbalances import RawTokenImbalances
 from src.helpers.helper_functions import get_web3_instance
 
 load_dotenv()
@@ -54,6 +54,6 @@ def test_imbalances(tx_hash, expected_imbalances):
     """
     chain_name = os.getenv("CHAIN_NAME")
     rt = RawTokenImbalances(get_web3_instance(), chain_name)
-    imbalances = rt.compute_imbalances(tx_hash)
+    imbalances = rt.compute_token_imbalances(tx_hash)
     for token_address, expected_imbalance in expected_imbalances.items():
         assert imbalances.get(token_address) == expected_imbalance

tests/legacy/compare_imbalances.py

@@ -6,7 +6,7 @@
 import time
 from web3 import Web3
 from src.helpers.config import ETHEREUM_NODE_URL
-from src.imbalances_script import RawTokenImbalances
+from src.raw_imbalances import RawTokenImbalances
 from src.balanceof_imbalances import BalanceOfImbalances
 from src.daemon import get_web3_instance, create_db_connection, fetch_transaction_hashes