Skip to main content

reth_e2e_test_utils/
payload.rs

1use futures_util::StreamExt;
2use reth_node_api::{PayloadAttributes, PayloadKind};
3use reth_payload_builder::{PayloadBuilderHandle, PayloadId};
4use reth_payload_builder_primitives::Events;
5use reth_payload_primitives::PayloadTypes;
6use tokio_stream::wrappers::BroadcastStream;
7
8/// Helper for payload operations
9#[derive(derive_more::Debug)]
10pub struct PayloadTestContext<T: PayloadTypes> {
11    pub payload_event_stream: BroadcastStream<Events<T>>,
12    payload_builder: PayloadBuilderHandle<T>,
13    /// Timestamp of the last generated payload attributes, or the starting point if none were
14    /// generated yet.
15    ///
16    /// [`Self::next_attributes`] increments it before generating the next attributes. Prefer
17    /// [`NodeTestContext::set_next_payload_timestamp`] over setting it directly.
18    ///
19    /// [`NodeTestContext::set_next_payload_timestamp`]: crate::node::NodeTestContext::set_next_payload_timestamp
20    pub timestamp: u64,
21    #[debug(skip)]
22    attributes_generator: Box<dyn Fn(u64) -> T::PayloadAttributes + Send + Sync>,
23}
24
25impl<T: PayloadTypes> PayloadTestContext<T> {
26    /// Creates a new payload helper
27    pub async fn new(
28        payload_builder: PayloadBuilderHandle<T>,
29        attributes_generator: impl Fn(u64) -> T::PayloadAttributes + Send + Sync + 'static,
30    ) -> eyre::Result<Self> {
31        let payload_events = payload_builder.subscribe().await?;
32        let payload_event_stream = payload_events.into_stream();
33        // Cancun timestamp
34        Ok(Self {
35            payload_event_stream,
36            payload_builder,
37            timestamp: 1710338135,
38            attributes_generator: Box::new(attributes_generator),
39        })
40    }
41
42    /// Generates the next payload attributes
43    pub fn next_attributes(&mut self) -> T::PayloadAttributes {
44        self.timestamp += 1;
45        (self.attributes_generator)(self.timestamp)
46    }
47
48    /// Asserts that the next event is a payload attributes event
49    pub async fn expect_attr_event(&mut self, attrs: T::PayloadAttributes) -> eyre::Result<()> {
50        let first_event = self.payload_event_stream.next().await.unwrap()?;
51        if let Events::Attributes(attr) = first_event {
52            assert_eq!(attrs.timestamp(), attr.timestamp());
53        } else {
54            panic!("Expect first event as payload attributes.")
55        }
56        Ok(())
57    }
58
59    /// Wait until the best built payload is ready.
60    ///
61    /// Panics if the payload builder does not produce a non-empty payload within 30 seconds.
62    pub async fn wait_for_built_payload(&self, payload_id: PayloadId) {
63        self.payload_builder
64            .resolve_kind(payload_id, PayloadKind::WaitForPending)
65            .await
66            .unwrap()
67            .unwrap();
68    }
69
70    /// Expects the next event to be a built payload event or panics
71    pub async fn expect_built_payload(&mut self) -> eyre::Result<T::BuiltPayload> {
72        let second_event = self.payload_event_stream.next().await.unwrap()?;
73        if let Events::BuiltPayload(payload) = second_event {
74            Ok(payload)
75        } else {
76            panic!("Expect a built payload event.");
77        }
78    }
79}