feat(testing): add balances check

testing/evm.py

@@ -1,24 +1,36 @@
+import os
 from web3 import Web3
 
+native_aliases = ["0x0000000000000000000000000000000000000000","0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee"]
+
+erc20_abi = [
+    {
+        "constant": True,
+        "inputs": [{"name": "_owner", "type": "address"}],
+        "name": "balanceOf",
+        "outputs": [{"name": "balance", "type": "uint256"}],
+        "type": "function",
+    }
+]
+
+def get_token_balance(token_address, wallet_address, block_number):
+    rpc_url = os.getenv("RPC_URL")
+
+    if rpc_url is None:
+        raise EnvironmentError("RPC_URL environment variable not set")
 
-def get_token_balance(rpc_url, token_address, wallet_address, block_number):
     web3 = Web3(Web3.HTTPProvider(rpc_url))
 
     if not web3.isConnected():
         raise ConnectionError("Failed to connect to the Ethereum node")
-
-    erc20_abi = [
-        {
-            "constant": True,
-            "inputs": [{"name": "_owner", "type": "address"}],
-            "name": "balanceOf",
-            "outputs": [{"name": "balance", "type": "uint256"}],
-            "type": "function",
-        }
-    ]
-
-    contract = web3.eth.contract(address=token_address, abi=erc20_abi)
-    balance = contract.functions.balanceOf(wallet_address).call(
-        block_identifier=block_number
-    )
+    
+    # Check if the token_address is a native token alias
+    if token_address.lower() in native_aliases:
+        balance = web3.eth.get_balance(Web3.toChecksumAddress(wallet_address), block_identifier=block_number)
+    else:
+        contract = web3.eth.contract(address=Web3.toChecksumAddress(token_address), abi=erc20_abi)
+        balance = contract.functions.balanceOf(Web3.toChecksumAddress(wallet_address)).call(
+            block_identifier=block_number
+        )
+    
     return balance

testing/runner.py

@@ -5,6 +5,7 @@ import subprocess
 
 import yaml
 
+from evm import get_token_balance
 from tycho import TychoRunner
 
 
@@ -51,7 +52,7 @@ class TestRunner:
             )
 
             result = self.tycho_runner.run_with_rpc_server(
-                self.validate_state, test["expected_state"]
+                self.validate_state, test["expected_state"], test["stop_block"]
             )
 
             if result.success:
@@ -63,9 +64,10 @@ class TestRunner:
                 "postgres://postgres:mypassword@localhost:5432"
             )
 
-    def validate_state(self, expected_state: dict) -> TestResult:
+    def validate_state(self, expected_state: dict, stop_block: int) -> TestResult:
         """Validate the current protocol state against the expected state."""
         protocol_components = self.tycho_runner.get_protocol_components()
+        protocol_states = self.tycho_runner.get_protocol_state()
         components = {
             component["id"]: component
             for component in protocol_components["protocol_components"]
@@ -96,6 +98,12 @@ class TestRunner:
                         return TestResult.Failed(
                             f"Value mismatch for key '{key}': {value} != {component[key]}"
                         )
+                for state in protocol_states["states"]:
+                    for token, balance in state["balances"].items():
+                        node_balance = get_token_balance(token,comp_id,stop_block)
+                        tycho_balance = int(balance,16)
+                        if node_balance != tycho_balance:
+                            return TestResult.Failed(f"Balance mismatch for {comp_id}:{token} at block {stop_block}: got {node_balance} from rpc call and {tycho_balance} from Substreams")
             return TestResult.Passed()
 
         except Exception as e:

testing/tycho.py

@@ -141,10 +141,7 @@ class TychoRunner:
         """Retrieve protocol state from the RPC server."""
         url = "http://0.0.0.0:4242/v1/ethereum/protocol_state"
         headers = {"accept": "application/json", "Content-Type": "application/json"}
-        data = {
-            "protocolSystem": "string",
-            "version": {"block": {"chain": "ethereum", "number": 0}},
-        }
+        data = {}
 
         response = requests.post(url, headers=headers, json=data)
         return response.json()