diff --git a/protocol-testing/src/rpc.rs b/protocol-testing/src/rpc.rs index 1709923..7ad7a33 100644 --- a/protocol-testing/src/rpc.rs +++ b/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 { 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)); } diff --git a/protocol-testing/src/test_runner.rs b/protocol-testing/src/test_runner.rs index 07f1c5a..6a30a5e 100644 --- a/protocol-testing/src/test_runner.rs +++ b/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, + protocol_states_by_id: &HashMap, + 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" + ); + } } diff --git a/substreams/ethereum-ekubo-v2/integration_test.tycho.yaml b/substreams/ethereum-ekubo-v2/integration_test.tycho.yaml index 91a9126..1faaa9e 100644 --- a/substreams/ethereum-ekubo-v2/integration_test.tycho.yaml +++ b/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"