feat: build snapshot message from rpc requests
This commit is contained in:
@@ -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() };
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user