From 71a05c3f47e98052921b1e709438af8a3731d422 Mon Sep 17 00:00:00 2001 From: Omar Abdulla Date: Mon, 11 Aug 2025 14:31:06 +0300 Subject: [PATCH] Add a balance assertion test step --- crates/core/src/driver/mod.rs | 84 +++++++++++++++++++++++++++++- crates/core/src/main.rs | 1 + crates/format/src/input.rs | 16 ++++++ crates/node-interaction/src/lib.rs | 4 ++ crates/node/src/geth.rs | 9 ++++ crates/node/src/kitchensink.rs | 9 ++++ 6 files changed, 122 insertions(+), 1 deletion(-) diff --git a/crates/core/src/driver/mod.rs b/crates/core/src/driver/mod.rs index e0b6bae..1f72021 100644 --- a/crates/core/src/driver/mod.rs +++ b/crates/core/src/driver/mod.rs @@ -26,7 +26,9 @@ use revive_dt_format::traits::{ResolutionContext, ResolverApi}; use semver::Version; use revive_dt_format::case::{Case, CaseIdx}; -use revive_dt_format::input::{Calldata, EtherValue, Expected, ExpectedOutput, Input, Method}; +use revive_dt_format::input::{ + BalanceAssertion, Calldata, EtherValue, Expected, ExpectedOutput, Input, Method, +}; use revive_dt_format::metadata::{ContractInstance, ContractPathAndIdent}; use revive_dt_format::{input::Step, metadata::Metadata}; use revive_dt_node::Node; @@ -83,6 +85,11 @@ where self.handle_input(metadata, case_idx, input, node).await?; Ok(StepOutput::FunctionCall(receipt, geth_trace, diff_mode)) } + Step::BalanceAssertion(balance_assertion) => { + self.handle_balance_assertion(metadata, case_idx, balance_assertion, node) + .await?; + Ok(StepOutput::BalanceAssertion) + } } } @@ -109,6 +116,19 @@ where .await } + pub async fn handle_balance_assertion( + &mut self, + metadata: &Metadata, + _: CaseIdx, + balance_assertion: &BalanceAssertion, + node: &T::Blockchain, + ) -> anyhow::Result<()> { + self.handle_balance_assertion_contract_deployment(metadata, balance_assertion, node) + .await?; + + Ok(()) + } + /// Handles the contract deployment for a given input performing it if it needs to be performed. async fn handle_input_contract_deployment( &mut self, @@ -478,6 +498,65 @@ where Ok((execution_receipt, trace, diff)) } + pub async fn handle_balance_assertion_contract_deployment( + &mut self, + metadata: &Metadata, + balance_assertion: &BalanceAssertion, + node: &T::Blockchain, + ) -> anyhow::Result<()> { + let Some(instance) = balance_assertion + .address + .strip_prefix(".address") + .map(ContractInstance::new) + else { + return Ok(()); + }; + self.get_or_deploy_contract_instance( + &instance, + metadata, + Input::default_caller(), + None, + None, + node, + ) + .await?; + Ok(()) + } + + pub async fn handle_balance_assertion_execution( + &mut self, + BalanceAssertion { + address: address_string, + amount, + }: &BalanceAssertion, + node: &T::Blockchain, + ) -> anyhow::Result<()> { + let address = Address::from_slice( + Calldata::new_compound([address_string]) + .calldata(node, self.default_resolution_context()) + .await? + .get(12..32) + .expect("Can't fail"), + ); + + let balance = node.balance_of(address).await?; + + let expected = *amount; + let actual = balance; + if expected != actual { + tracing::error!(%expected, %actual, %address, "Balance assertion failed"); + anyhow::bail!( + "Balance assertion failed - Expected {} but got {} for {} resolved to {}", + expected, + actual, + address_string, + address, + ) + } + + Ok(()) + } + /// Gets the information of a deployed contract or library from the state. If it's found to not /// be deployed then it will be deployed. /// @@ -700,6 +779,8 @@ where tracing::trace!("Follower logs: {:?}", follower_receipt.logs()); } } + (StepOutput::BalanceAssertion, StepOutput::BalanceAssertion) => {} + _ => unreachable!("The two step outputs can not be of a different kind"), } steps_executed += 1; @@ -712,4 +793,5 @@ where #[derive(Clone, Debug)] pub enum StepOutput { FunctionCall(TransactionReceipt, GethTrace, DiffMode), + BalanceAssertion, } diff --git a/crates/core/src/main.rs b/crates/core/src/main.rs index 1a87cbc..70512a5 100644 --- a/crates/core/src/main.rs +++ b/crates/core/src/main.rs @@ -450,6 +450,7 @@ where .iter() .filter_map(|step| match step { Step::FunctionCall(input) => Some(input.caller), + Step::BalanceAssertion(..) => None, }) .next() .unwrap_or(Input::default_caller()); diff --git a/crates/format/src/input.rs b/crates/format/src/input.rs index 5f95d07..437b5d5 100644 --- a/crates/format/src/input.rs +++ b/crates/format/src/input.rs @@ -26,6 +26,8 @@ use crate::{metadata::ContractInstance, traits::ResolutionContext}; pub enum Step { /// A function call or an invocation to some function on some smart contract. FunctionCall(Input), + /// A step for performing a balance assertion on some account or contract. + BalanceAssertion(BalanceAssertion), } #[derive(Clone, Debug, Default, Serialize, Deserialize, Eq, PartialEq)] @@ -44,6 +46,20 @@ pub struct Input { pub variable_assignments: Option, } +#[derive(Clone, Debug, Default, Serialize, Deserialize, Eq, PartialEq)] +pub struct BalanceAssertion { + /// The address that the balance assertion should be done on. + /// + /// This is a string which will be resolved into an address when being processed. Therefore, + /// this could be a normal hex address, a variable such as `Test.address`, or perhaps even a + /// full on variable like `$VARIABLE:Uniswap`. It follows the same resolution rules that are + /// followed in the calldata. + pub address: String, + + /// The amount of balance to assert that the account or contract has. + pub amount: U256, +} + #[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)] #[serde(untagged)] pub enum Expected { diff --git a/crates/node-interaction/src/lib.rs b/crates/node-interaction/src/lib.rs index 791ba4b..b052c7a 100644 --- a/crates/node-interaction/src/lib.rs +++ b/crates/node-interaction/src/lib.rs @@ -1,5 +1,6 @@ //! This crate implements all node interactions. +use alloy::primitives::{Address, U256}; use alloy::rpc::types::trace::geth::{DiffMode, GethDebugTracingOptions, GethTrace}; use alloy::rpc::types::{TransactionReceipt, TransactionRequest}; use anyhow::Result; @@ -21,4 +22,7 @@ pub trait EthereumNode { /// Returns the state diff of the transaction hash in the [TransactionReceipt]. fn state_diff(&self, receipt: &TransactionReceipt) -> impl Future>; + + /// Returns the balance of the provided [`Address`] back. + fn balance_of(&self, address: Address) -> impl Future>; } diff --git a/crates/node/src/geth.rs b/crates/node/src/geth.rs index 5ce7921..b3618ba 100644 --- a/crates/node/src/geth.rs +++ b/crates/node/src/geth.rs @@ -371,6 +371,15 @@ impl EthereumNode for GethNode { _ => anyhow::bail!("expected a diff mode trace"), } } + + #[tracing::instrument(skip_all, fields(geth_node_id = self.id))] + async fn balance_of(&self, address: Address) -> anyhow::Result { + self.provider() + .await? + .get_balance(address) + .await + .map_err(Into::into) + } } impl ResolverApi for GethNode { diff --git a/crates/node/src/kitchensink.rs b/crates/node/src/kitchensink.rs index 2300e05..805b71f 100644 --- a/crates/node/src/kitchensink.rs +++ b/crates/node/src/kitchensink.rs @@ -428,6 +428,15 @@ impl EthereumNode for KitchensinkNode { _ => anyhow::bail!("expected a diff mode trace"), } } + + #[tracing::instrument(skip_all, fields(kitchensink_node_id = self.id))] + async fn balance_of(&self, address: Address) -> anyhow::Result { + self.provider() + .await? + .get_balance(address) + .await + .map_err(Into::into) + } } impl ResolverApi for KitchensinkNode {