diff --git a/testing/evm.py b/testing/evm.py index 4408140..1d90989 100644 --- a/testing/evm.py +++ b/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 diff --git a/testing/runner.py b/testing/runner.py index 2189570..deab2e3 100644 --- a/testing/runner.py +++ b/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: diff --git a/testing/tycho.py b/testing/tycho.py index 988dc70..1806b0a 100644 --- a/testing/tycho.py +++ b/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()