From 9f4ff71b7072d38d535c3143514245d39a8cc459 Mon Sep 17 00:00:00 2001
From: adrian <adri.benavides312@gmail.com>
Date: Thu, 18 Sep 2025 10:10:02 +0200
Subject: [PATCH] refactor: remove `skip_balance_check` from run_test signature

This field is redundant, as it can be read directly from `config`
---
 protocol-testing/src/test_runner.rs  | 16 +++++++---------
 protocol-testing/src/tycho_runner.rs |  5 ++---
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/protocol-testing/src/test_runner.rs b/protocol-testing/src/test_runner.rs
index c8faa29..e859fcc 100644
--- a/protocol-testing/src/test_runner.rs
+++ b/protocol-testing/src/test_runner.rs
@@ -121,7 +121,7 @@ impl TestRunner {
         for test in &tests {
             info!("TEST {}: {}", count, test.name);
 
-            match self.run_test(test, &config, config.skip_balance_check) {
+            match self.run_test(test, &config) {
                 Ok(_) => {
                     info!("✅ {} passed\n", test.name);
                 }
@@ -159,7 +159,6 @@ impl TestRunner {
         &self,
         test: &IntegrationTest,
         config: &IntegrationTestsConfig,
-        skip_balance_check: bool,
     ) -> miette::Result<()> {
         self.empty_database()
             .into_diagnostic()
@@ -195,12 +194,11 @@ impl TestRunner {
             .wrap_err("Failed to run Tycho")?;
 
         tycho_runner.run_with_rpc_server(
-            |expected_components, start_block, stop_block, skip_balance_check| {
+            |expected_components, start_block, stop_block| {
                 validate_state(
                     expected_components,
                     start_block,
                     stop_block,
-                    skip_balance_check,
                     config,
                     &self.adapter_contract_builder,
                     self.vm_traces,
@@ -209,7 +207,6 @@ impl TestRunner {
             &test.expected_components,
             test.start_block,
             test.stop_block,
-            skip_balance_check,
         )?
     }
 
@@ -233,7 +230,6 @@ fn validate_state(
     expected_components: &Vec<ProtocolComponentWithTestConfig>,
     start_block: u64,
     stop_block: u64,
-    skip_balance_check: bool,
     config: &IntegrationTestsConfig,
     adapter_contract_builder: &AdapterContractBuilder,
     vm_traces: bool,
@@ -290,7 +286,10 @@ fn validate_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();
+        let component_id = expected_component
+            .base
+            .id
+            .to_lowercase();
 
         let component = components_by_id
             .get(&component_id)
@@ -315,7 +314,7 @@ fn validate_state(
     info!("All expected components were successfully found on Tycho and match the expected state");
 
     // Step 2: Validate Token Balances
-    match skip_balance_check {
+    match config.skip_balance_check {
         true => info!("Skipping balance check"),
         false => {
             validate_token_balances(&components_by_id, &protocol_states_by_id, start_block, &rt)?;
@@ -545,7 +544,6 @@ fn validate_state(
             }
         }
     }
-
     info!("✅ Simulation validation passed");
     Ok(())
 }
diff --git a/protocol-testing/src/tycho_runner.rs b/protocol-testing/src/tycho_runner.rs
index 26dab37..c35eeed 100644
--- a/protocol-testing/src/tycho_runner.rs
+++ b/protocol-testing/src/tycho_runner.rs
@@ -98,10 +98,9 @@ impl TychoRunner {
         expected_components: &Vec<ProtocolComponentWithTestConfig>,
         start_block: u64,
         stop_block: u64,
-        skip_balance_check: bool,
     ) -> miette::Result<R>
     where
-        F: FnOnce(&Vec<ProtocolComponentWithTestConfig>, u64, u64, bool) -> R,
+        F: FnOnce(&Vec<ProtocolComponentWithTestConfig>, u64, u64) -> R,
     {
         let (tx, rx): (Sender<bool>, Receiver<bool>) = mpsc::channel();
         let db_url = self.db_url.clone();
@@ -141,7 +140,7 @@ impl TychoRunner {
         thread::sleep(Duration::from_secs(3));
 
         // Run the provided function
-        let result = func(expected_components, start_block, stop_block, skip_balance_check);
+        let result = func(expected_components, start_block, stop_block);
 
         tx.send(true)
             .expect("Failed to send termination message");