feat: build snapshot message from rpc requests

This commit is contained in:
Thales
2025-03-22 14:41:42 -03:00
committed by Tamara
parent 052d11b759
commit c9b2c655dd
5 changed files with 1048 additions and 299 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -6,11 +6,10 @@ use alloy::{
providers::{Provider, ProviderBuilder},
transports::http::reqwest::Url,
};
use tycho_core::models::Address;
const NATIVE_ALIASES: &[Address] = &[
address!("0x0000000000000000000000000000000000000000").into(),
address!("0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee").into(),
address!("0x0000000000000000000000000000000000000000"),
address!("0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee"),
];
const ERC_20_ABI: &str = r#"[{"inputs":[{"name":"_owner","type":"address"}],"name":"balanceOf","outputs":[{"name":"balance","type":"uint256"}],"stateMutability":"view","type":"function"}]"#;
@@ -20,18 +19,19 @@ pub struct RPCProvider {
}
impl RPCProvider {
pub(crate) fn new(url: String) -> RPCProvider {
pub fn new(url: String) -> Self {
let url = url.as_str().parse().unwrap();
RPCProvider { url }
}
// TODO: Return a Result instead of panicking
pub async fn get_token_balance(
self,
&self,
token_address: Address,
wallet_address: Address,
block_number: u64,
) -> U256 {
let provider = ProviderBuilder::new().on_http(self.url);
let provider = ProviderBuilder::new().on_http(self.url.clone());
let block_id: BlockId = BlockId::from(block_number);
match NATIVE_ALIASES.contains(&token_address) {
@@ -65,7 +65,7 @@ impl RPCProvider {
}
}
async fn get_block_header(self, block_number: u64) {
async fn get_block_header(&self, block_number: u64) {
// TODO: Implement
// let provider = ProviderBuilder::new().on_http(self.url);
// let block_id: BlockId = BlockId::from(block_number);
@@ -84,7 +84,7 @@ mod tests {
#[tokio::test]
async fn get_token_balance_native_token() {
let eth_rpc_url = env::var("ETH_RPC_URL").expect("Missing ETH_RPC_URL in environment");
let eth_rpc_url = env::var("RPC_URL").expect("Missing RPC_URL in environment");
let rpc_provider = RPCProvider::new(eth_rpc_url);
let token_address = address!("0x0000000000000000000000000000000000000000");
@@ -103,7 +103,7 @@ mod tests {
#[tokio::test]
async fn get_token_balance_erc20_token() {
let eth_rpc_url = env::var("ETH_RPC_URL").expect("Missing ETH_RPC_URL in environment");
let eth_rpc_url = env::var("RPC_URL").expect("Missing RPC_URL in environment");
let rpc_provider = RPCProvider::new(eth_rpc_url);
let token_address = address!("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48");

View File

@@ -1,4 +1,4 @@
use std::{collections::HashMap, env, ops::Deref, path::PathBuf};
use std::{collections::HashMap, env, ops::Deref, path::PathBuf, str::FromStr};
use alloy::{
primitives::{bytes, U256},
@@ -12,19 +12,29 @@ use postgres::{Client, Error, NoTls};
use tokio::runtime::Runtime;
use tracing::{debug, field::debug, info};
use tycho_core::{
dto::{Chain, ProtocolComponent, ResponseProtocolState},
dto::{Chain, ProtocolComponent, ResponseAccount, ResponseProtocolState},
models::Address,
Bytes,
};
use tycho_core::models::Address;
use tycho_simulation::evm::protocol::u256_num::{bytes_to_u256, u256_to_f64};
use tycho_simulation::{
evm::{
decoder::TychoStreamDecoder,
engine_db::tycho_db::PreCachedDB,
protocol::{
u256_num::{bytes_to_u256, u256_to_f64},
vm::state::EVMPoolState,
},
},
tycho_client::feed::{synchronizer::StateSyncMessage, FeedMessage, Header},
};
use tycho_simulation::tycho_client::feed::synchronizer::{ComponentW, ComponentWithState, Snapshot};
use crate::{
config::{IntegrationTest, IntegrationTestsConfig, ProtocolComponentWithTestConfig},
rpc::RPCProvider,
tycho_rpc::TychoClient,
tycho_runner::TychoRunner,
utils::build_spkg,
};
use crate::rpc::RPCProvider;
pub struct TestRunner {
package: String,
@@ -57,7 +67,7 @@ impl TestRunner {
info!("Found {} tests to run", config.tests.len());
for test in &config.tests {
self.run_test(test, &config);
self.run_test(test, &config, config.skip_balance_check);
}
}
Err(e) => {
@@ -66,7 +76,12 @@ impl TestRunner {
}
}
fn run_test(&self, test: &IntegrationTest, config: &IntegrationTestsConfig, skip_balance_check: bool) {
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");
@@ -101,7 +116,13 @@ impl TestRunner {
)
.expect("Failed to run Tycho");
tycho_runner.run_with_rpc_server(validate_state, &test.expected_components, test.start_block, skip_balance_check);
tycho_runner.run_with_rpc_server(
validate_state,
&test.expected_components,
test.start_block,
test.stop_block,
skip_balance_check,
);
}
fn empty_database(&self) -> Result<(), Error> {
@@ -122,7 +143,12 @@ impl TestRunner {
}
}
fn validate_state(expected_components: &Vec<ProtocolComponentWithTestConfig>, start_block: u64, skip_balance_check: bool) {
fn validate_state(
expected_components: &Vec<ProtocolComponentWithTestConfig>,
start_block: u64,
stop_block: u64,
skip_balance_check: bool,
) {
let rt = Runtime::new().unwrap();
// Create Tycho client for the RPC server
@@ -132,6 +158,8 @@ fn validate_state(expected_components: &Vec<ProtocolComponentWithTestConfig>, st
let chain = Chain::Ethereum;
let protocol_system = "test_protocol";
// Fetch data from Tycho RPC. We use block_on to avoid using async functions on the testing
// module, in order to simplify debugging
let protocol_components = rt
.block_on(tycho_client.get_protocol_components(protocol_system, chain))
.expect("Failed to get protocol components");
@@ -140,15 +168,20 @@ fn validate_state(expected_components: &Vec<ProtocolComponentWithTestConfig>, st
.block_on(tycho_client.get_protocol_state(protocol_system, chain))
.expect("Failed to get protocol state");
let vm_storages = rt
.block_on(tycho_client.get_contract_state(Vec::new(), protocol_system, chain))
.expect("Failed to get contract state");
// Create a map of component IDs to components for easy lookup
let components_by_id: HashMap<String, ProtocolComponent> = protocol_components
.clone()
.into_iter()
.map(|c| (c.id.to_lowercase(), c))
.map(|c| (c.id.clone(), c))
.collect();
let protocol_states_by_id: HashMap<String, ResponseProtocolState> = protocol_states
.into_iter()
.map(|s| (s.component_id.to_lowercase(), s))
.map(|s| (s.component_id.clone(), s))
.collect();
info!("Found {} protocol components", components_by_id.len());
@@ -159,10 +192,7 @@ fn validate_state(expected_components: &Vec<ProtocolComponentWithTestConfig>, st
// 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.clone();
assert!(
components_by_id.contains_key(&component_id),
@@ -194,8 +224,8 @@ fn validate_state(expected_components: &Vec<ProtocolComponentWithTestConfig>, st
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 (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);
@@ -208,10 +238,73 @@ fn validate_state(expected_components: &Vec<ProtocolComponentWithTestConfig>, st
}
}
// TODO: Test if balance check works
if (!skip_balance_check) {
let token_address: Address
let node_balance = rpc_provider.get_token_balance(token, component.id, start_block)
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,
));
assert_eq!(
balance, node_balance,
"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"),
}
// Step 3: Run Tycho Simulation
let mut decoder = TychoStreamDecoder::new();
decoder.register_decoder::<EVMPoolState<PreCachedDB>>("test_protocol");
// Mock a stream message, with only a Snapshot and no deltas
let mut states: HashMap<String, ComponentWithState> = HashMap::new();
for (id, component) in components_by_id {
let component_id = &id.clone();
let state = protocol_states_by_id
.get(component_id)
.expect("Failed to get state for component")
.clone();
let component_with_state = ComponentWithState { state, component };
states.insert(component_id.clone(), component_with_state);
}
let vm_storage: HashMap<Bytes, ResponseAccount> = vm_storages
.into_iter()
.map(|x| (x.address.clone(), x))
.collect();
let snapshot = Snapshot { states, vm_storage };
let state_msgs: HashMap<String, StateSyncMessage> = HashMap::from([(
String::from("test_protocol"),
StateSyncMessage {
header: Header {
hash: Default::default(),
number: stop_block,
parent_hash: Default::default(),
revert: false,
},
snapshots: snapshot,
deltas: None,
removed_components: HashMap::new(),
},
)]);
let stream_message: FeedMessage = FeedMessage { state_msgs, sync_states: Default::default() };
}

View File

@@ -99,10 +99,11 @@ impl TychoRunner {
func: F,
expected_components: &Vec<ProtocolComponentWithTestConfig>,
start_block: u64,
stop_block: u64,
skip_balance_check: bool,
) -> R
where
F: FnOnce(&Vec<ProtocolComponentWithTestConfig>, u64, bool) -> R,
F: FnOnce(&Vec<ProtocolComponentWithTestConfig>, u64, u64, bool) -> R,
{
let (tx, rx): (Sender<bool>, Receiver<bool>) = mpsc::channel();
let db_url = self.db_url.clone();
@@ -141,7 +142,7 @@ impl TychoRunner {
thread::sleep(Duration::from_secs(3));
// Run the provided function
let result = func(expected_components, start_block, skip_balance_check);
let result = func(expected_components, start_block, stop_block, skip_balance_check);
tx.send(true)
.expect("Failed to send termination message");

View File

@@ -1,4 +1,5 @@
use std::{
env,
error::Error,
fs,
path::{Path, PathBuf},
@@ -10,6 +11,7 @@ use figment::{
value::Value,
Figment,
};
use tracing::error;
/// Build a Substreams package with modifications to the YAML file.
pub fn build_spkg(yaml_file_path: &PathBuf, initial_block: u64) -> Result<String, Box<dyn Error>> {
@@ -55,6 +57,7 @@ pub fn build_spkg(yaml_file_path: &PathBuf, initial_block: u64) -> Result<String
fs::write(yaml_file_path, yaml_string)?;
// Run the substreams pack command to create the spkg
// WARNING: Ensure substreams is in the PATH
match Command::new("substreams")
.arg("pack")
.arg(yaml_file_path)
@@ -62,14 +65,16 @@ pub fn build_spkg(yaml_file_path: &PathBuf, initial_block: u64) -> Result<String
{
Ok(output) => {
if !output.status.success() {
println!(
error!(
"Substreams pack command failed: {}",
String::from_utf8_lossy(&output.stderr)
);
}
}
Err(e) => {
println!("Error running substreams pack command: {}", e);
error!("Error running substreams pack command: {}. \
Ensure that the wasm target was built and that substreams CLI\
is installed and exported on PATH", e);
}
}