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