Add a balance assertion test step

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,
 }

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());

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<VariableAssignments>,
 }
 
+#[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 {

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<Output = Result<DiffMode>>;
+
+    /// Returns the balance of the provided [`Address`] back.
+    fn balance_of(&self, address: Address) -> impl Future<Output = Result<U256>>;
 }

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<U256> {
+        self.provider()
+            .await?
+            .get_balance(address)
+            .await
+            .map_err(Into::into)
+    }
 }
 
 impl ResolverApi for GethNode {

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<U256> {
+        self.provider()
+            .await?
+            .get_balance(address)
+            .await
+            .map_err(Into::into)
+    }
 }
 
 impl ResolverApi for KitchensinkNode {