test(exex): introduce utils for testing (#8636)

Cargo.lock

@@ -2808,9 +2808,11 @@ dependencies = [
  "futures",
  "reth",
  "reth-exex",
+ "reth-exex-test-utils",
  "reth-node-api",
  "reth-node-ethereum",
  "reth-tracing",
+ "tokio",
 ]
 
 [[package]]
@@ -6912,6 +6914,32 @@ dependencies = [
  "tokio-util",
 ]
 
+[[package]]
+name = "reth-exex-test-utils"
+version = "0.2.0-beta.8"
+dependencies = [
+ "eyre",
+ "futures-util",
+ "rand 0.8.5",
+ "reth-blockchain-tree",
+ "reth-config",
+ "reth-db",
+ "reth-db-common",
+ "reth-evm",
+ "reth-exex",
+ "reth-network",
+ "reth-node-api",
+ "reth-node-builder",
+ "reth-node-core",
+ "reth-node-ethereum",
+ "reth-payload-builder",
+ "reth-primitives",
+ "reth-provider",
+ "reth-tasks",
+ "reth-transaction-pool",
+ "tokio",
+]
+
 [[package]]
 name = "reth-fs-util"
 version = "0.2.0-beta.8"

Cargo.toml

@@ -21,7 +21,8 @@ members = [
     "crates/evm/",
     "crates/evm/execution-errors",
     "crates/evm/execution-types",
-    "crates/exex/",
+    "crates/exex/exex/",
+    "crates/exex/test-utils/",
     "crates/metrics/",
     "crates/metrics/metrics-derive/",
     "crates/net/common/",
@@ -263,7 +264,8 @@ reth-evm-ethereum = { path = "crates/ethereum/evm" }
 reth-evm-optimism = { path = "crates/optimism/evm" }
 reth-execution-errors = { path = "crates/evm/execution-errors" }
 reth-execution-types = { path = "crates/evm/execution-types" }
-reth-exex = { path = "crates/exex" }
+reth-exex = { path = "crates/exex/exex" }
+reth-exex-test-utils = { path = "crates/exex/test-utils" }
 reth-fs-util = { path = "crates/fs-util" }
 reth-ipc = { path = "crates/rpc/ipc" }
 reth-libmdbx = { path = "crates/storage/libmdbx-rs" }

crates/exex/exex/src/context.rs

@@ -1,3 +1,5 @@
+use std::fmt::Debug;
+
 use reth_node_api::{FullNodeComponents, FullNodeTypes, NodeTypes};
 use reth_node_core::node_config::NodeConfig;
 use reth_primitives::Head;
@@ -7,7 +9,6 @@ use tokio::sync::mpsc::{Receiver, UnboundedSender};
 use crate::{ExExEvent, ExExNotification};
 
 /// Captures the context that an `ExEx` has access to.
-#[derive(Debug)]
 pub struct ExExContext<Node: FullNodeComponents> {
     /// The current head of the blockchain at launch.
     pub head: Head,
@@ -35,6 +36,19 @@ pub struct ExExContext<Node: FullNodeComponents> {
     pub components: Node,
 }
 
+impl<Node: FullNodeComponents> Debug for ExExContext<Node> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ExExContext")
+            .field("head", &self.head)
+            .field("config", &self.config)
+            .field("reth_config", &self.reth_config)
+            .field("events", &self.events)
+            .field("notifications", &self.notifications)
+            .field("components", &"...")
+            .finish()
+    }
+}
+
 impl<Node: FullNodeComponents> NodeTypes for ExExContext<Node> {
     type Primitives = Node::Primitives;
     type Engine = Node::Engine;

crates/exex/test-utils/Cargo.toml

@@ -0,0 +1,38 @@
+[package]
+name = "reth-exex-test-utils"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+## reth
+reth-blockchain-tree.workspace = true
+reth-config.workspace = true
+reth-db = { workspace = true, features = ["test-utils"] }
+reth-db-common.workspace = true
+reth-evm = { workspace = true, features = ["test-utils"] }
+reth-exex.workspace = true
+reth-network.workspace = true
+reth-node-api.workspace = true
+reth-node-core.workspace = true
+reth-node-builder.workspace = true
+reth-node-ethereum.workspace = true
+reth-payload-builder.workspace = true
+reth-primitives.workspace = true
+reth-provider.workspace = true
+reth-tasks.workspace = true
+reth-transaction-pool = { workspace = true, features = ["test-utils"] }
+
+## async
+futures-util.workspace = true
+tokio.workspace = true
+
+## misc
+eyre.workspace = true
+rand.workspace = true

crates/exex/test-utils/src/lib.rs

@@ -0,0 +1,260 @@
+//! Test helpers for `reth-exex`
+
+#![doc(
+    html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png",
+    html_favicon_url = "https://avatars0.githubusercontent.com/u/97369466?s=256",
+    issue_tracker_base_url = "https://github.com/paradigmxyz/reth/issues/"
+)]
+#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
+#![cfg_attr(not(test), warn(unused_crate_dependencies))]
+
+use futures_util::FutureExt;
+use reth_blockchain_tree::noop::NoopBlockchainTree;
+use reth_db::{test_utils::TempDatabase, DatabaseEnv};
+use reth_db_common::init::init_genesis;
+use reth_evm::test_utils::MockExecutorProvider;
+use reth_exex::{ExExContext, ExExEvent, ExExNotification};
+use reth_network::{config::SecretKey, NetworkConfigBuilder, NetworkManager};
+use reth_node_api::{FullNodeTypes, FullNodeTypesAdapter, NodeTypes};
+use reth_node_builder::{
+    components::{
+        Components, ComponentsBuilder, ExecutorBuilder, NodeComponentsBuilder, PoolBuilder,
+    },
+    BuilderContext, Node, NodeAdapter, RethFullAdapter,
+};
+use reth_node_core::node_config::NodeConfig;
+use reth_node_ethereum::{
+    node::{EthereumNetworkBuilder, EthereumPayloadBuilder},
+    EthEngineTypes, EthEvmConfig,
+};
+use reth_payload_builder::noop::NoopPayloadBuilderService;
+use reth_primitives::{ChainSpec, Head, SealedBlockWithSenders, MAINNET};
+use reth_provider::{
+    providers::BlockchainProvider, test_utils::create_test_provider_factory_with_chain_spec,
+    BlockReader, Chain, ProviderFactory,
+};
+use reth_tasks::TaskManager;
+use reth_transaction_pool::test_utils::{testing_pool, TestPool};
+use std::{
+    fmt::Debug,
+    future::{poll_fn, Future},
+    sync::Arc,
+    task::Poll,
+};
+use tokio::sync::mpsc::{Sender, UnboundedReceiver};
+
+/// A test [`PoolBuilder`] that builds a [`TestPool`].
+#[derive(Debug, Default, Clone, Copy)]
+#[non_exhaustive]
+pub struct TestPoolBuilder;
+
+impl<Node> PoolBuilder<Node> for TestPoolBuilder
+where
+    Node: FullNodeTypes,
+{
+    type Pool = TestPool;
+
+    async fn build_pool(self, _ctx: &BuilderContext<Node>) -> eyre::Result<Self::Pool> {
+        Ok(testing_pool())
+    }
+}
+
+/// A test [`ExecutorBuilder`] that builds a [`MockExecutorProvider`].
+#[derive(Debug, Default, Clone, Copy)]
+#[non_exhaustive]
+pub struct TestExecutorBuilder;
+
+impl<Node> ExecutorBuilder<Node> for TestExecutorBuilder
+where
+    Node: FullNodeTypes,
+{
+    type EVM = EthEvmConfig;
+    type Executor = MockExecutorProvider;
+
+    async fn build_evm(
+        self,
+        _ctx: &BuilderContext<Node>,
+    ) -> eyre::Result<(Self::EVM, Self::Executor)> {
+        let evm_config = EthEvmConfig::default();
+        let executor = MockExecutorProvider::default();
+
+        Ok((evm_config, executor))
+    }
+}
+
+/// A test [`Node`].
+#[derive(Debug, Default, Clone, Copy)]
+#[non_exhaustive]
+pub struct TestNode;
+
+impl NodeTypes for TestNode {
+    type Primitives = ();
+    type Engine = EthEngineTypes;
+}
+
+impl<N> Node<N> for TestNode
+where
+    N: FullNodeTypes<Engine = EthEngineTypes>,
+{
+    type ComponentsBuilder = ComponentsBuilder<
+        N,
+        TestPoolBuilder,
+        EthereumPayloadBuilder,
+        EthereumNetworkBuilder,
+        TestExecutorBuilder,
+    >;
+
+    fn components_builder(self) -> Self::ComponentsBuilder {
+        ComponentsBuilder::default()
+            .node_types::<N>()
+            .pool(TestPoolBuilder::default())
+            .payload(EthereumPayloadBuilder::default())
+            .network(EthereumNetworkBuilder::default())
+            .executor(TestExecutorBuilder::default())
+    }
+}
+
+type TmpDB = Arc<TempDatabase<DatabaseEnv>>;
+type Adapter = NodeAdapter<
+    RethFullAdapter<TmpDB, TestNode>,
+    <<TestNode as Node<FullNodeTypesAdapter<TestNode, TmpDB, BlockchainProvider<TmpDB>>>>::ComponentsBuilder as NodeComponentsBuilder<
+        RethFullAdapter<TmpDB, TestNode>,
+    >>::Components,
+>;
+
+/// A helper type for testing Execution Extensions.
+#[derive(Debug)]
+pub struct TestExExHandle {
+    /// Genesis block that was inserted into the storage
+    pub genesis: SealedBlockWithSenders,
+    /// Provider Factory for accessing the emphemeral storage of the host node
+    pub provider_factory: ProviderFactory<TmpDB>,
+    /// Channel for receiving events from the Execution Extension
+    pub events_rx: UnboundedReceiver<ExExEvent>,
+    /// Channel for sending notifications to the Execution Extension
+    pub notifications_tx: Sender<ExExNotification>,
+}
+
+impl TestExExHandle {
+    /// Send a notification to the Execution Extension that the chain has been committed
+    pub async fn send_notification_chain_committed(&self, chain: Chain) -> eyre::Result<()> {
+        self.notifications_tx
+            .send(ExExNotification::ChainCommitted { new: Arc::new(chain) })
+            .await?;
+        Ok(())
+    }
+
+    /// Asserts that the Execution Extension did not emit any events.
+    #[track_caller]
+    pub fn assert_events_empty(&self) {
+        assert!(self.events_rx.is_empty());
+    }
+
+    /// Asserts that the Execution Extension emitted a `FinishedHeight` event with the correct
+    /// height.
+    #[track_caller]
+    pub fn assert_event_finished_height(&mut self, height: u64) -> eyre::Result<()> {
+        let event = self.events_rx.try_recv()?;
+        assert_eq!(event, ExExEvent::FinishedHeight(height));
+        Ok(())
+    }
+}
+
+/// Creates a new [`ExExContext`].
+///
+/// This is a convenience function that does the following:
+/// 1. Sets up an [`ExExContext`] with all dependencies.
+/// 2. Inserts the genesis block of the provided (chain spec)[`ChainSpec`] into the storage.
+/// 3. Creates a channel for receiving events from the Execution Extension.
+/// 4. Creates a channel for sending notifications to the Execution Extension.
+///
+/// # Warning
+/// The genesis block is not sent to the notifications channel. The caller is responsible for
+/// doing this.
+pub async fn test_exex_context_with_chain_spec(
+    chain_spec: Arc<ChainSpec>,
+) -> eyre::Result<(ExExContext<Adapter>, TestExExHandle)> {
+    let transaction_pool = testing_pool();
+    let evm_config = EthEvmConfig::default();
+    let executor = MockExecutorProvider::default();
+
+    let provider_factory = create_test_provider_factory_with_chain_spec(chain_spec);
+    let genesis_hash = init_genesis(provider_factory.clone())?;
+    let provider =
+        BlockchainProvider::new(provider_factory.clone(), Arc::new(NoopBlockchainTree::default()))?;
+
+    let network_manager = NetworkManager::new(
+        NetworkConfigBuilder::new(SecretKey::new(&mut rand::thread_rng()))
+            .build(provider_factory.clone()),
+    )
+    .await?;
+    let network = network_manager.handle().clone();
+
+    let (_, payload_builder) = NoopPayloadBuilderService::<EthEngineTypes>::new();
+
+    let tasks = TaskManager::current();
+    let task_executor = tasks.executor();
+
+    let components = NodeAdapter::<FullNodeTypesAdapter<TestNode, _, _>, _> {
+        components: Components { transaction_pool, evm_config, executor, network, payload_builder },
+        task_executor,
+        provider,
+    };
+
+    let genesis = provider_factory
+        .block_by_hash(genesis_hash)?
+        .ok_or(eyre::eyre!("genesis block not found"))?
+        .seal_slow()
+        .seal_with_senders()
+        .ok_or(eyre::eyre!("failed to recover senders"))?;
+
+    let head = Head {
+        number: genesis.number,
+        hash: genesis_hash,
+        difficulty: genesis.difficulty,
+        timestamp: genesis.timestamp,
+        total_difficulty: Default::default(),
+    };
+
+    let (events_tx, events_rx) = tokio::sync::mpsc::unbounded_channel();
+    let (notifications_tx, notifications_rx) = tokio::sync::mpsc::channel(1);
+
+    let ctx = ExExContext {
+        head,
+        config: NodeConfig::test(),
+        reth_config: reth_config::Config::default(),
+        events: events_tx,
+        notifications: notifications_rx,
+        components,
+    };
+
+    Ok((ctx, TestExExHandle { genesis, provider_factory, events_rx, notifications_tx }))
+}
+
+/// Creates a new [`ExExContext`] with (mainnet)[`MAINNET`] chain spec.
+///
+/// For more information see [`test_exex_context_with_chain_spec`].
+pub async fn test_exex_context() -> eyre::Result<(ExExContext<Adapter>, TestExExHandle)> {
+    test_exex_context_with_chain_spec(MAINNET.clone()).await
+}
+
+/// An extension trait for polling an Execution Extension future.
+pub trait PollOnce {
+    /// Polls the given Execution Extension future __once__ and asserts that it is
+    /// [`Poll::Pending`]. The future should be (pinned)[`std::pin::pin`].
+    ///
+    /// # Panics
+    /// If the future returns [`Poll::Ready`], because Execution Extension future should never
+    /// resolve.
+    fn poll_once(&mut self) -> impl Future<Output = ()> + Send;
+}
+
+impl<F: Future<Output = eyre::Result<()>> + Unpin + Send> PollOnce for F {
+    async fn poll_once(&mut self) {
+        poll_fn(|cx| {
+            assert!(self.poll_unpin(cx).is_pending());
+            Poll::Ready(())
+        })
+        .await;
+    }
+}

examples/exex/minimal/Cargo.toml

@@ -14,3 +14,8 @@ reth-tracing.workspace = true
 
 eyre.workspace = true
 futures.workspace = true
+
+[dev-dependencies]
+reth-exex-test-utils.workspace = true
+
+tokio.workspace = true

examples/exex/minimal/src/main.rs

@@ -36,6 +36,7 @@ async fn exex<Node: FullNodeComponents>(mut ctx: ExExContext<Node>) -> eyre::Res
             ctx.events.send(ExExEvent::FinishedHeight(committed_chain.tip().number))?;
         }
     }
+
     Ok(())
 }
 
@@ -50,3 +51,44 @@ fn main() -> eyre::Result<()> {
         handle.wait_for_node_exit().await
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use std::pin::pin;
+
+    use reth::providers::{BundleStateWithReceipts, Chain};
+    use reth_exex_test_utils::{test_exex_context, PollOnce};
+
+    #[tokio::test]
+    async fn exex() -> eyre::Result<()> {
+        // Initialize a test Execution Extension context with all dependencies
+        let (ctx, mut handle) = test_exex_context().await?;
+
+        // Save the current head of the chain to check the finished height against it later
+        let head = ctx.head;
+
+        // Send a notification to the Execution Extension that the chain has been committed
+        handle
+            .send_notification_chain_committed(Chain::from_block(
+                handle.genesis.clone(),
+                BundleStateWithReceipts::default(),
+                None,
+            ))
+            .await?;
+
+        // Initialize the Execution Extension
+        let mut exex = pin!(super::exex_init(ctx).await?);
+
+        // Check that the Execution Extension did not emit any events until we polled it
+        handle.assert_events_empty();
+
+        // Poll the Execution Extension once to process incoming notifications
+        exex.poll_once().await;
+
+        // Check that the Execution Extension emitted a `FinishedHeight` event with the correct
+        // height
+        handle.assert_event_finished_height(head.number)?;
+
+        Ok(())
+    }
+}