test: add test for the token balances validation
This commit is contained in:
@@ -6,6 +6,7 @@ use alloy::{
|
||||
providers::{Provider, ProviderBuilder},
|
||||
transports::http::reqwest::Url,
|
||||
};
|
||||
use miette::{IntoDiagnostic, WrapErr};
|
||||
|
||||
const NATIVE_ALIASES: &[Address] = &[
|
||||
address!("0x0000000000000000000000000000000000000000"),
|
||||
@@ -24,13 +25,12 @@ impl RPCProvider {
|
||||
RPCProvider { url }
|
||||
}
|
||||
|
||||
// TODO: Return a Result instead of panicking
|
||||
pub async fn get_token_balance(
|
||||
&self,
|
||||
token_address: Address,
|
||||
wallet_address: Address,
|
||||
block_number: u64,
|
||||
) -> U256 {
|
||||
) -> miette::Result<U256> {
|
||||
let provider = ProviderBuilder::new().connect_http(self.url.clone());
|
||||
let block_id: BlockId = BlockId::from(block_number);
|
||||
|
||||
@@ -39,9 +39,12 @@ impl RPCProvider {
|
||||
.get_balance(wallet_address)
|
||||
.block_id(block_id)
|
||||
.await
|
||||
.expect("Failed to fetch token balance"),
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to fetch token balance"),
|
||||
false => {
|
||||
let abi = serde_json::from_str(ERC_20_ABI).expect("invalid ABI");
|
||||
let abi = serde_json::from_str(ERC_20_ABI)
|
||||
.into_diagnostic()
|
||||
.wrap_err("invalid ABI")?;
|
||||
|
||||
let contract = ContractInstance::new(token_address, provider, Interface::new(abi));
|
||||
|
||||
@@ -53,14 +56,15 @@ impl RPCProvider {
|
||||
.block(block_id)
|
||||
.call()
|
||||
.await
|
||||
.expect("Failed to fetch ERC-20 Balance");
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to fetch ERC-20 Balance")?;
|
||||
let result: U256 = result_value
|
||||
.first()
|
||||
.unwrap()
|
||||
.ok_or_else(|| miette::miette!("No value returned from contract call"))?
|
||||
.as_uint()
|
||||
.unwrap()
|
||||
.ok_or_else(|| miette::miette!("Returned value is not a uint"))?
|
||||
.0;
|
||||
result
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -93,7 +97,8 @@ mod tests {
|
||||
|
||||
let balance = rpc_provider
|
||||
.get_token_balance(token_address, wallet_address, block_number)
|
||||
.await;
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
balance,
|
||||
@@ -112,7 +117,8 @@ mod tests {
|
||||
|
||||
let balance = rpc_provider
|
||||
.get_token_balance(token_address, wallet_address, block_number)
|
||||
.await;
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(balance, U256::from(717250938432_u64));
|
||||
}
|
||||
|
||||
@@ -254,58 +254,12 @@ fn validate_state(
|
||||
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")
|
||||
.into_diagnostic()
|
||||
.wrap_err("Missing ETH_RPC_URL in environment")?;
|
||||
let rpc_provider = RPCProvider::new(rpc_url.to_string());
|
||||
|
||||
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);
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: Test if balance check works
|
||||
if !skip_balance_check {
|
||||
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,
|
||||
));
|
||||
if balance != node_balance {
|
||||
return Err(miette!(
|
||||
"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"),
|
||||
false => {
|
||||
validate_token_balances(&components_by_id, &protocol_states_by_id, start_block, &rt)?;
|
||||
info!("All token balances match the values found onchain")
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3: Run Tycho Simulation
|
||||
@@ -404,9 +358,68 @@ fn validate_state(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validate that the token balances of the components match the values
|
||||
/// on-chain, extracted by querying the token balances using a node.
|
||||
fn validate_token_balances(
|
||||
components_by_id: &HashMap<String, ProtocolComponent>,
|
||||
protocol_states_by_id: &HashMap<String, ResponseProtocolState>,
|
||||
start_block: u64,
|
||||
rt: &Runtime,
|
||||
) -> miette::Result<()> {
|
||||
let rpc_url = env::var("RPC_URL")
|
||||
.into_diagnostic()
|
||||
.wrap_err("Missing RPC_URL in environment")?;
|
||||
let rpc_provider = RPCProvider::new(rpc_url.to_string());
|
||||
|
||||
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);
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
))?;
|
||||
if balance != node_balance {
|
||||
return Err(miette!(
|
||||
"Token balance mismatch for component {} and token {}",
|
||||
component.id,
|
||||
token
|
||||
));
|
||||
}
|
||||
info!(
|
||||
"Token balance for component {} and token {} matches the expected value",
|
||||
component.id, token
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use dotenv::dotenv;
|
||||
use glob::glob;
|
||||
use tycho_common::{
|
||||
dto::{ProtocolComponent, ResponseProtocolState},
|
||||
Bytes,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -455,4 +468,80 @@ mod tests {
|
||||
panic!("One or more config files failed to parse.");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_balance_validation() {
|
||||
// Setup mock data
|
||||
let block_number = 21998530;
|
||||
let token_bytes = Bytes::from_str("0x0000000000000000000000000000000000000000").unwrap();
|
||||
let component_id = "0x787B8840100d9BaAdD7463f4a73b5BA73B00C6cA".to_string();
|
||||
|
||||
let component = ProtocolComponent {
|
||||
id: component_id.clone(),
|
||||
tokens: vec![token_bytes.clone()],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut balances = HashMap::new();
|
||||
let balance_bytes = Bytes::from(
|
||||
U256::from_str("1070041574684539264153")
|
||||
.unwrap()
|
||||
.to_be_bytes::<32>(),
|
||||
);
|
||||
balances.insert(token_bytes.clone(), balance_bytes.clone());
|
||||
let protocol_state = ResponseProtocolState {
|
||||
component_id: component_id.clone(),
|
||||
balances,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut components_by_id = HashMap::new();
|
||||
components_by_id.insert(component_id.clone(), component.clone());
|
||||
let mut protocol_states_by_id = HashMap::new();
|
||||
protocol_states_by_id.insert(component_id.clone(), protocol_state.clone());
|
||||
|
||||
let rt = Runtime::new().unwrap();
|
||||
dotenv().ok();
|
||||
let result =
|
||||
validate_token_balances(&components_by_id, &protocol_states_by_id, block_number, &rt);
|
||||
assert!(result.is_ok(), "Should pass when balance check is performed and balances match");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_balance_validation_fails_on_mismatch() {
|
||||
// Setup mock data
|
||||
let block_number = 21998530;
|
||||
let token_bytes = Bytes::from_str("0x0000000000000000000000000000000000000000").unwrap();
|
||||
let component_id = "0x787B8840100d9BaAdD7463f4a73b5BA73B00C6cA".to_string();
|
||||
|
||||
let component = ProtocolComponent {
|
||||
id: component_id.clone(),
|
||||
tokens: vec![token_bytes.clone()],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Set expected balance to zero
|
||||
let mut balances = HashMap::new();
|
||||
let balance_bytes = Bytes::from(U256::from(0).to_be_bytes::<32>());
|
||||
balances.insert(token_bytes.clone(), balance_bytes.clone());
|
||||
let protocol_state = ResponseProtocolState {
|
||||
component_id: component_id.clone(),
|
||||
balances,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut components_by_id = HashMap::new();
|
||||
components_by_id.insert(component_id.clone(), component.clone());
|
||||
let mut protocol_states_by_id = HashMap::new();
|
||||
protocol_states_by_id.insert(component_id.clone(), protocol_state.clone());
|
||||
|
||||
let rt = Runtime::new().unwrap();
|
||||
dotenv().ok();
|
||||
let result =
|
||||
validate_token_balances(&components_by_id, &protocol_states_by_id, block_number, &rt);
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Should fail when balance check is performed and balances do not match"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
substreams_yaml_path: ./substreams.yaml
|
||||
adapter_contract: "EkuboSwapAdapter"
|
||||
adapter_build_signature: "constructor(address)"
|
||||
adapter_build_args: "0x16e186ecdc94083fff53ef2a41d46b92a54f61e2"
|
||||
|
||||
Reference in New Issue
Block a user