feat: Add tokens balance check and component diff

protocol-testing/src/config.rs

@@ -1,36 +1,19 @@
 use std::collections::HashMap;
 
+use colored::Colorize;
 use hex::FromHex;
 use serde::{Deserialize, Serialize};
-
-/// Represents a hexadecimal byte string: Check if we already have a default impl for this
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
-#[serde(from = "String", into = "String")]
-pub struct HexBytes(Vec<u8>);
-
-impl From<String> for HexBytes {
-    fn from(s: String) -> Self {
-        let s = s
-            .trim_start_matches("0x")
-            .to_lowercase();
-        HexBytes(Vec::from_hex(s).unwrap_or_default())
-    }
-}
-
-impl From<HexBytes> for String {
-    fn from(val: HexBytes) -> Self {
-        format!("0x{}", hex::encode(val.0))
-    }
-}
+use similar::{ChangeTag, TextDiff};
+use tycho_core::{dto::ProtocolComponent, Bytes};
 
 /// Represents a ProtocolComponent with its main attributes
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct ProtocolComponentExpectation {
     pub id: String,
-    pub tokens: Vec<HexBytes>,
+    pub tokens: Vec<Bytes>,
     #[serde(default)]
-    pub static_attributes: HashMap<String, HexBytes>,
-    pub creation_tx: HexBytes,
+    pub static_attributes: HashMap<String, Bytes>,
+    pub creation_tx: Bytes,
 }
 
 /// Represents a ProtocolComponent with test configuration
@@ -42,6 +25,98 @@ pub struct ProtocolComponentWithTestConfig {
     pub skip_simulation: bool,
 }
 
+impl ProtocolComponentExpectation {
+    pub fn compare(&self, other: &ProtocolComponent, colorize_output: bool) -> Option<String> {
+        let mut diffs = Vec::new();
+
+        // Compare id
+        if self.id != other.id {
+            let diff = self.format_diff("id", &self.id, &other.id, colorize_output);
+            diffs.push(format!("Field 'id' mismatch for {}:\n{}", self.id, diff));
+        }
+
+        // Compare tokens
+        if self.tokens != other.tokens {
+            let self_tokens = format!("{:?}", self.tokens);
+            let other_tokens = format!("{:?}", other.tokens);
+            let diff = self.format_diff("tokens", &self_tokens, &other_tokens, colorize_output);
+            diffs.push(format!("Field 'tokens' mismatch for {}:\n{}", self.id, diff));
+        }
+
+        // Compare static_attributes
+        for (key, value) in &self.static_attributes {
+            let other_value = other.static_attributes.get(key);
+            match other_value {
+                Some(other_value) => {
+                    if value != other_value {
+                        let self_value = format!("{:?}", value);
+                        let other_value = format!("{:?}", other_value);
+                        let diff = self.format_diff(
+                            "static_attributes",
+                            &self_value,
+                            &other_value,
+                            colorize_output,
+                        );
+                        diffs.push(format!(
+                            "Field 'static_attributes' mismatch for {}:\n{}",
+                            self.id, diff
+                        ));
+                    }
+                }
+                None => {
+                    diffs.push(format!(
+                        "Field 'static_attributes' mismatch for {}: Key '{}' not found",
+                        self.id, key
+                    ));
+                }
+            }
+        }
+        // Compare creation_tx
+        if self.creation_tx != other.creation_tx {
+            let self_tx = format!("{}", self.creation_tx.clone());
+            let other_tx = format!("{}", other.creation_tx.clone());
+            let diff = self.format_diff("creation_tx", &self_tx, &other_tx, colorize_output);
+            diffs.push(format!("Field 'creation_tx' mismatch for {}:\n{}", self.id, diff));
+        }
+
+        if diffs.is_empty() {
+            None
+        } else {
+            Some(diffs.join("\n"))
+        }
+    }
+    fn format_diff(&self, _field_name: &str, left: &str, right: &str, colorize: bool) -> String {
+        let diff = TextDiff::from_lines(left, right);
+
+        let mut result = String::new();
+        for change in diff.iter_all_changes() {
+            let formatted = match change.tag() {
+                ChangeTag::Delete => {
+                    if colorize {
+                        format!("{}", format!("-{}", change.value().trim_end()).red())
+                    } else {
+                        format!("-{}", change.value().trim_end())
+                    }
+                }
+                ChangeTag::Insert => {
+                    if colorize {
+                        format!("{}", format!("+{}", change.value().trim_end()).green())
+                    } else {
+                        format!("+{}", change.value().trim_end())
+                    }
+                }
+                ChangeTag::Equal => {
+                    format!(" {}", change.value().trim_end())
+                }
+            };
+            result.push_str(&formatted);
+            result.push('\n');
+        }
+
+        result
+    }
+}
+
 fn default_false() -> bool {
     false
 }

protocol-testing/src/rpc.rs

@@ -6,25 +6,26 @@ use alloy::{
     providers::{Provider, ProviderBuilder},
     transports::http::reqwest::Url,
 };
+use tycho_core::models::Address;
 
 const NATIVE_ALIASES: &[Address] = &[
-    address!("0x0000000000000000000000000000000000000000"),
-    address!("0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee"),
+    address!("0x0000000000000000000000000000000000000000").into(),
+    address!("0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee").into(),
 ];
 
 const ERC_20_ABI: &str = r#"[{"inputs":[{"name":"_owner","type":"address"}],"name":"balanceOf","outputs":[{"name":"balance","type":"uint256"}],"stateMutability":"view","type":"function"}]"#;
 
-struct RPCProvider {
+pub struct RPCProvider {
     url: Url,
 }
 
 impl RPCProvider {
-    fn new(url: String) -> RPCProvider {
+    pub(crate) fn new(url: String) -> RPCProvider {
         let url = url.as_str().parse().unwrap();
         RPCProvider { url }
     }
 
-    async fn get_token_balance(
+    pub async fn get_token_balance(
         self,
         token_address: Address,
         wallet_address: Address,

protocol-testing/src/test_runner.rs

@@ -1,20 +1,30 @@
-use std::{collections::HashMap, path::PathBuf};
+use std::{collections::HashMap, env, ops::Deref, path::PathBuf};
 
+use alloy::{
+    primitives::{bytes, U256},
+    providers::Provider,
+};
 use figment::{
     providers::{Format, Yaml},
     Figment,
 };
 use postgres::{Client, Error, NoTls};
 use tokio::runtime::Runtime;
-use tracing::{debug, info};
-use tycho_core::dto::{Chain, ProtocolComponent};
+use tracing::{debug, field::debug, info};
+use tycho_core::{
+    dto::{Chain, ProtocolComponent, ResponseProtocolState},
+    Bytes,
+};
+use tycho_core::models::Address;
+use tycho_simulation::evm::protocol::u256_num::{bytes_to_u256, u256_to_f64};
 
 use crate::{
-    config::{IntegrationTest, IntegrationTestsConfig},
+    config::{IntegrationTest, IntegrationTestsConfig, ProtocolComponentWithTestConfig},
     tycho_rpc::TychoClient,
     tycho_runner::TychoRunner,
     utils::build_spkg,
 };
+use crate::rpc::RPCProvider;
 
 pub struct TestRunner {
     package: String,
@@ -56,7 +66,7 @@ impl TestRunner {
         }
     }
 
-    fn run_test(&self, test: &IntegrationTest, config: &IntegrationTestsConfig) {
+    fn run_test(&self, test: &IntegrationTest, config: &IntegrationTestsConfig, skip_balance_check: bool) {
         info!("Running test: {}", test.name);
         self.empty_database()
             .expect("Failed to empty the database");
@@ -91,7 +101,7 @@ impl TestRunner {
             )
             .expect("Failed to run Tycho");
 
-        tycho_runner.run_with_rpc_server(validate_state);
+        tycho_runner.run_with_rpc_server(validate_state, &test.expected_components, test.start_block, skip_balance_check);
     }
 
     fn empty_database(&self) -> Result<(), Error> {
@@ -112,7 +122,7 @@ impl TestRunner {
     }
 }
 
-fn validate_state() {
+fn validate_state(expected_components: &Vec<ProtocolComponentWithTestConfig>, start_block: u64, skip_balance_check: bool) {
     let rt = Runtime::new().unwrap();
 
     // Create Tycho client for the RPC server
@@ -133,12 +143,75 @@ fn validate_state() {
     // Create a map of component IDs to components for easy lookup
     let components_by_id: HashMap<String, ProtocolComponent> = protocol_components
         .into_iter()
-        .map(|c| (c.id.clone(), c))
+        .map(|c| (c.id.to_lowercase(), c))
+        .collect();
+
+    let protocol_states_by_id: HashMap<String, ResponseProtocolState> = protocol_states
+        .into_iter()
+        .map(|s| (s.component_id.to_lowercase(), s))
         .collect();
 
     info!("Found {} protocol components", components_by_id.len());
-    info!("Found {} protocol states", protocol_states.len());
+    info!("Found {} protocol states", protocol_states_by_id.len());
 
-    // TODO: Implement complete validation logic similar to Python code
     info!("Validating state...");
+
+    // Step 1: Validate that all expected components are present on Tycho after indexing
+    debug!("Validating {:?} expected components", expected_components.len());
+    for expected_component in expected_components {
+        let component_id = expected_component
+            .base
+            .id
+            .to_lowercase();
+
+        assert!(
+            components_by_id.contains_key(&component_id),
+            "Component {:?} was not found on Tycho",
+            component_id
+        );
+
+        let component = components_by_id
+            .get(&component_id)
+            .unwrap();
+
+        let diff = expected_component
+            .base
+            .compare(&component, true);
+        match diff {
+            Some(diff) => {
+                panic!("Component {} does not match the expected state:\n{}", component_id, diff);
+            }
+            None => {
+                info!("Component {} matches the expected state", component_id);
+            }
+        }
+    }
+    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").expect("Missing ETH_RPC_URL in environment");
+    let rpc_provider = RPCProvider::new(rpc_url.to_string());
+
+    for (id_lower, component) in components_by_id.iter() {
+        let component_state = protocol_states_by_id.get(id_lower);
+
+        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);
+                }
+            }
+
+            if (!skip_balance_check) {
+                let token_address: Address
+                let node_balance = rpc_provider.get_token_balance(token, component.id, start_block)
+            }
+        }
+    }
 }

protocol-testing/src/tycho_runner.rs

@@ -9,6 +9,8 @@ use std::{
 use dotenv::dotenv;
 use tracing::debug;
 
+use crate::config::ProtocolComponentWithTestConfig;
+
 pub struct TychoRunner {
     db_url: String,
     initialized_accounts: Vec<String>,
@@ -92,9 +94,15 @@ impl TychoRunner {
         Ok(())
     }
 
-    pub fn run_with_rpc_server<F, R>(&self, func: F) -> R
+    pub fn run_with_rpc_server<F, R>(
+        &self,
+        func: F,
+        expected_components: &Vec<ProtocolComponentWithTestConfig>,
+        start_block: u64,
+        skip_balance_check: bool,
+    ) -> R
     where
-        F: FnOnce() -> R,
+        F: FnOnce(&Vec<ProtocolComponentWithTestConfig>, u64, bool) -> R,
     {
         let (tx, rx): (Sender<bool>, Receiver<bool>) = mpsc::channel();
         let db_url = self.db_url.clone();
@@ -133,7 +141,7 @@ impl TychoRunner {
         thread::sleep(Duration::from_secs(3));
 
         // Run the provided function
-        let result = func();
+        let result = func(expected_components, start_block, skip_balance_check);
 
         tx.send(true)
             .expect("Failed to send termination message");