From 160523a8887c89c25306d29ae6793d77493647ac Mon Sep 17 00:00:00 2001
From: Thales Lima <thales@datarevenue.com>
Date: Tue, 11 Mar 2025 13:59:02 -0300
Subject: [PATCH] feat: add test runner structure

---
 protocol-testing/src/main.rs         |  43 ++++++++++
 protocol-testing/src/test_runner.rs  | 119 +++++++++++++++++++++++++++
 protocol-testing/src/tycho_runner.rs |   2 +-
 3 files changed, 163 insertions(+), 1 deletion(-)
 create mode 100644 protocol-testing/src/main.rs
 create mode 100644 protocol-testing/src/test_runner.rs

diff --git a/protocol-testing/src/main.rs b/protocol-testing/src/main.rs
new file mode 100644
index 0000000..a347328
--- /dev/null
+++ b/protocol-testing/src/main.rs
@@ -0,0 +1,43 @@
+mod config;
+mod rpc;
+mod test_runner;
+mod utils;
+mod tycho_runner;
+
+use clap::Parser;
+use tracing_subscriber::EnvFilter;
+
+use crate::test_runner::TestRunner;
+
+#[derive(Parser, Debug)]
+#[command(version, about = "Run indexer within a specified range of blocks")]
+struct Args {
+    /// Name of the package to test
+    #[arg(long)]
+    package: String,
+
+    /// Enable tycho logs
+    #[arg(long, default_value_t = true)]
+    tycho_logs: bool,
+
+    /// Postgres database URL for the Tycho indexer
+    #[arg(long, default_value = "postgres://postgres:mypassword@localhost:5431/tycho_indexer_0")]
+    db_url: String,
+
+    /// Enable tracing during vm simulations
+    #[arg(long, default_value_t = false)]
+    vm_traces: bool,
+}
+
+fn main() {
+    tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env())
+        .with_target(false)
+        .init();
+
+    let args = Args::parse();
+
+    let test_runner = TestRunner::new(args.package, args.tycho_logs, args.db_url, args.vm_traces);
+
+    test_runner.run_tests();
+}
diff --git a/protocol-testing/src/test_runner.rs b/protocol-testing/src/test_runner.rs
new file mode 100644
index 0000000..5fb2853
--- /dev/null
+++ b/protocol-testing/src/test_runner.rs
@@ -0,0 +1,119 @@
+use std::{
+    path::{Path, PathBuf},
+    thread::sleep,
+    time::Duration,
+};
+
+use figment::{
+    providers::{Format, Yaml},
+    Figment,
+};
+use postgres::{Client, Error, NoTls};
+use tracing::{debug, field::debug, info};
+
+use crate::{
+    config::{IntegrationTest, IntegrationTestsConfig},
+    tycho_runner::TychoRunner,
+    utils::{build_spkg, modify_initial_block},
+};
+
+pub struct TestRunner {
+    package: String,
+    tycho_logs: bool,
+    db_url: String,
+    vm_traces: bool,
+
+    substreams_path: PathBuf,
+}
+
+impl TestRunner {
+    pub fn new(package: String, tycho_logs: bool, db_url: String, vm_traces: bool) -> Self {
+        let substreams_path = PathBuf::from("../substreams").join(&package);
+        Self { package, tycho_logs, db_url, vm_traces, substreams_path }
+    }
+
+    pub fn run_tests(&self) {
+        info!("Running tests...");
+        let config_yaml_path = self
+            .substreams_path
+            .join("integration_test.tycho.yaml");
+
+        info!("Config YAML: {}", config_yaml_path.display());
+        let figment = Figment::new().merge(Yaml::file(&config_yaml_path));
+
+        match figment.extract::<IntegrationTestsConfig>() {
+            Ok(config) => {
+                info!("Loaded test configuration:");
+                info!("Protocol types: {:?}", config.protocol_type_names);
+                info!("Found {} tests to run", config.tests.len());
+
+                for test in &config.tests {
+                    self.run_test(test, &config);
+                }
+            }
+            Err(e) => {
+                eprintln!("Failed to load test configuration: {}", e);
+            }
+        }
+    }
+
+    fn run_test(&self, test: &IntegrationTest, config: &IntegrationTestsConfig) {
+        info!("Running test: {}", test.name);
+        self.empty_database()
+            .expect("Failed to empty the database");
+
+        let substreams_yaml_path = self
+            .substreams_path
+            .join(&config.substreams_yaml_path);
+        debug!("Building SPKG on {:?}", substreams_yaml_path);
+
+        let mut initialized_accounts = config
+            .initialized_accounts
+            .clone()
+            .unwrap_or_default();
+        initialized_accounts.extend(
+            test.initialized_accounts
+                .clone()
+                .unwrap_or_default(),
+        );
+
+        let spkg_path =
+            build_spkg(&substreams_yaml_path, test.start_block).expect("Failed to build spkg");
+
+        let tycho_runner =
+            TychoRunner::new(self.db_url.clone(), initialized_accounts, self.tycho_logs);
+
+        tycho_runner
+            .run_tycho(
+                spkg_path.as_str(),
+                test.start_block,
+                test.stop_block,
+                &config.protocol_type_names,
+            )
+            .expect("Failed to run Tycho");
+
+        tycho_runner.run_with_rpc_server(validate_state)
+    }
+
+    fn empty_database(&self) -> Result<(), Error> {
+        debug!("Emptying the database");
+
+        // Remove db name from URL. This is required because we cannot drop a database that we are
+        // currently connected to.
+        let base_url = match self.db_url.rfind('/') {
+            Some(pos) => &self.db_url[..pos],
+            None => self.db_url.as_str(),
+        };
+        let mut client = Client::connect(base_url, NoTls)?;
+
+        client.execute("DROP DATABASE IF EXISTS \"tycho_indexer_0\" WITH (FORCE)", &[])?;
+        client.execute("CREATE DATABASE \"tycho_indexer_0\"", &[])?;
+
+        Ok(())
+    }
+}
+
+pub fn validate_state() {
+    // TODO: Implement
+    info!("Validating state...");
+}
diff --git a/protocol-testing/src/tycho_runner.rs b/protocol-testing/src/tycho_runner.rs
index 50e60f3..5ea8f0a 100644
--- a/protocol-testing/src/tycho_runner.rs
+++ b/protocol-testing/src/tycho_runner.rs
@@ -97,7 +97,7 @@ impl TychoRunner {
         Ok(())
     }
 
-    fn run_with_rpc_server<F, R>(&self, func: F) -> R
+    pub fn run_with_rpc_server<F, R>(&self, func: F) -> R
     where
         F: FnOnce() -> R,
     {