test: add test for the token balances validation

protocol-testing/src/rpc.rs

@@ -6,6 +6,7 @@ use alloy::{
     providers::{Provider, ProviderBuilder},
     transports::http::reqwest::Url,
 };
+use miette::{IntoDiagnostic, WrapErr};
 
 const NATIVE_ALIASES: &[Address] = &[
     address!("0x0000000000000000000000000000000000000000"),
@@ -24,13 +25,12 @@ impl RPCProvider {
         RPCProvider { url }
     }
 
-    // TODO: Return a Result instead of panicking
     pub async fn get_token_balance(
         &self,
         token_address: Address,
         wallet_address: Address,
         block_number: u64,
-    ) -> U256 {
+    ) -> miette::Result<U256> {
         let provider = ProviderBuilder::new().connect_http(self.url.clone());
         let block_id: BlockId = BlockId::from(block_number);
 
@@ -39,9 +39,12 @@ impl RPCProvider {
                 .get_balance(wallet_address)
                 .block_id(block_id)
                 .await
-                .expect("Failed to fetch token balance"),
+                .into_diagnostic()
+                .wrap_err("Failed to fetch token balance"),
             false => {
-                let abi = serde_json::from_str(ERC_20_ABI).expect("invalid ABI");
+                let abi = serde_json::from_str(ERC_20_ABI)
+                    .into_diagnostic()
+                    .wrap_err("invalid ABI")?;
 
                 let contract = ContractInstance::new(token_address, provider, Interface::new(abi));
 
@@ -53,14 +56,15 @@ impl RPCProvider {
                     .block(block_id)
                     .call()
                     .await
-                    .expect("Failed to fetch ERC-20 Balance");
+                    .into_diagnostic()
+                    .wrap_err("Failed to fetch ERC-20 Balance")?;
                 let result: U256 = result_value
                     .first()
-                    .unwrap()
+                    .ok_or_else(|| miette::miette!("No value returned from contract call"))?
                     .as_uint()
-                    .unwrap()
+                    .ok_or_else(|| miette::miette!("Returned value is not a uint"))?
                     .0;
-                result
+                Ok(result)
             }
         }
     }
@@ -93,7 +97,8 @@ mod tests {
 
         let balance = rpc_provider
             .get_token_balance(token_address, wallet_address, block_number)
-            .await;
+            .await
+            .unwrap();
 
         assert_eq!(
             balance,
@@ -112,7 +117,8 @@ mod tests {
 
         let balance = rpc_provider
             .get_token_balance(token_address, wallet_address, block_number)
-            .await;
+            .await
+            .unwrap();
 
         assert_eq!(balance, U256::from(717250938432_u64));
     }

protocol-testing/src/test_runner.rs

@@ -254,58 +254,12 @@ fn validate_state(
     info!("All expected components were successfully found on Tycho and match the expected state");
 
     // Step 2: Validate Token Balances
-    // In this step, we validate that the token balances of the components match the values
-    // on-chain, extracted by querying the token balances using a node.
-    let rpc_url = env::var("RPC_URL")
-        .into_diagnostic()
-        .wrap_err("Missing ETH_RPC_URL in environment")?;
-    let rpc_provider = RPCProvider::new(rpc_url.to_string());
-
-    for (id, component) in components_by_id.iter() {
-        let component_state = protocol_states_by_id.get(id);
-
-        for token in &component.tokens {
-            let mut balance: U256 = U256::from(0);
-
-            if let Some(state) = component_state {
-                let bal = state.balances.get(token);
-                if let Some(bal) = bal {
-                    let bal = bal.clone().into();
-                    balance = bytes_to_u256(bal);
-                }
-            }
-
-            // TODO: Test if balance check works
-            if !skip_balance_check {
-                info!(
-                    "Validating token balance for component {} and token {}",
-                    component.id, token
-                );
-                let token_address = alloy::primitives::Address::from_slice(&token[..20]);
-                let component_address = alloy::primitives::Address::from_str(component.id.as_str())
-                    .expect("Failed to parse component address");
-                let node_balance = rt.block_on(rpc_provider.get_token_balance(
-                    token_address,
-                    component_address,
-                    start_block,
-                ));
-                if balance != node_balance {
-                    return Err(miette!(
-                        "Token balance mismatch for component {} and token {}",
-                        component.id,
-                        token
-                    ))
-                }
-                info!(
-                    "Token balance for component {} and token {} matches the expected value",
-                    component.id, token
-                );
-            }
-        }
-    }
     match skip_balance_check {
         true => info!("Skipping balance check"),
-        false => info!("All token balances match the values found onchain"),
+        false => {
+            validate_token_balances(&components_by_id, &protocol_states_by_id, start_block, &rt)?;
+            info!("All token balances match the values found onchain")
+        }
     }
 
     // Step 3: Run Tycho Simulation
@@ -404,9 +358,68 @@ fn validate_state(
     Ok(())
 }
 
+/// Validate that the token balances of the components match the values
+/// on-chain, extracted by querying the token balances using a node.
+fn validate_token_balances(
+    components_by_id: &HashMap<String, ProtocolComponent>,
+    protocol_states_by_id: &HashMap<String, ResponseProtocolState>,
+    start_block: u64,
+    rt: &Runtime,
+) -> miette::Result<()> {
+    let rpc_url = env::var("RPC_URL")
+        .into_diagnostic()
+        .wrap_err("Missing RPC_URL in environment")?;
+    let rpc_provider = RPCProvider::new(rpc_url.to_string());
+
+    for (id, component) in components_by_id.iter() {
+        let component_state = protocol_states_by_id.get(id);
+
+        for token in &component.tokens {
+            let mut balance: U256 = U256::from(0);
+
+            if let Some(state) = component_state {
+                let bal = state.balances.get(token);
+                if let Some(bal) = bal {
+                    let bal = bal.clone().into();
+                    balance = bytes_to_u256(bal);
+                }
+            }
+
+            info!("Validating token balance for component {} and token {}", component.id, token);
+            let token_address = alloy::primitives::Address::from_slice(&token[..20]);
+            let component_address = alloy::primitives::Address::from_str(component.id.as_str())
+                .expect("Failed to parse component address");
+            let node_balance = rt.block_on(rpc_provider.get_token_balance(
+                token_address,
+                component_address,
+                start_block,
+            ))?;
+            if balance != node_balance {
+                return Err(miette!(
+                    "Token balance mismatch for component {} and token {}",
+                    component.id,
+                    token
+                ));
+            }
+            info!(
+                "Token balance for component {} and token {} matches the expected value",
+                component.id, token
+            );
+        }
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+
+    use dotenv::dotenv;
     use glob::glob;
+    use tycho_common::{
+        dto::{ProtocolComponent, ResponseProtocolState},
+        Bytes,
+    };
 
     use super::*;
 
@@ -455,4 +468,80 @@ mod tests {
             panic!("One or more config files failed to parse.");
         }
     }
+
+    #[test]
+    fn test_token_balance_validation() {
+        // Setup mock data
+        let block_number = 21998530;
+        let token_bytes = Bytes::from_str("0x0000000000000000000000000000000000000000").unwrap();
+        let component_id = "0x787B8840100d9BaAdD7463f4a73b5BA73B00C6cA".to_string();
+
+        let component = ProtocolComponent {
+            id: component_id.clone(),
+            tokens: vec![token_bytes.clone()],
+            ..Default::default()
+        };
+
+        let mut balances = HashMap::new();
+        let balance_bytes = Bytes::from(
+            U256::from_str("1070041574684539264153")
+                .unwrap()
+                .to_be_bytes::<32>(),
+        );
+        balances.insert(token_bytes.clone(), balance_bytes.clone());
+        let protocol_state = ResponseProtocolState {
+            component_id: component_id.clone(),
+            balances,
+            ..Default::default()
+        };
+
+        let mut components_by_id = HashMap::new();
+        components_by_id.insert(component_id.clone(), component.clone());
+        let mut protocol_states_by_id = HashMap::new();
+        protocol_states_by_id.insert(component_id.clone(), protocol_state.clone());
+
+        let rt = Runtime::new().unwrap();
+        dotenv().ok();
+        let result =
+            validate_token_balances(&components_by_id, &protocol_states_by_id, block_number, &rt);
+        assert!(result.is_ok(), "Should pass when balance check is performed and balances match");
+    }
+
+    #[test]
+    fn test_token_balance_validation_fails_on_mismatch() {
+        // Setup mock data
+        let block_number = 21998530;
+        let token_bytes = Bytes::from_str("0x0000000000000000000000000000000000000000").unwrap();
+        let component_id = "0x787B8840100d9BaAdD7463f4a73b5BA73B00C6cA".to_string();
+
+        let component = ProtocolComponent {
+            id: component_id.clone(),
+            tokens: vec![token_bytes.clone()],
+            ..Default::default()
+        };
+
+        // Set expected balance to zero
+        let mut balances = HashMap::new();
+        let balance_bytes = Bytes::from(U256::from(0).to_be_bytes::<32>());
+        balances.insert(token_bytes.clone(), balance_bytes.clone());
+        let protocol_state = ResponseProtocolState {
+            component_id: component_id.clone(),
+            balances,
+            ..Default::default()
+        };
+
+        let mut components_by_id = HashMap::new();
+        components_by_id.insert(component_id.clone(), component.clone());
+        let mut protocol_states_by_id = HashMap::new();
+        protocol_states_by_id.insert(component_id.clone(), protocol_state.clone());
+
+        let rt = Runtime::new().unwrap();
+        dotenv().ok();
+        let result =
+            validate_token_balances(&components_by_id, &protocol_states_by_id, block_number, &rt);
+        assert!(
+            result.is_err(),
+            "Should fail when balance check is performed and balances do not match"
+        );
+    }
 }

substreams/ethereum-ekubo-v2/integration_test.tycho.yaml

@@ -1,3 +1,4 @@
+substreams_yaml_path: ./substreams.yaml
 adapter_contract: "EkuboSwapAdapter"
 adapter_build_signature: "constructor(address)"
 adapter_build_args: "0x16e186ecdc94083fff53ef2a41d46b92a54f61e2"