reth_e2e_test_utils/
payload.rs

1use futures_util::StreamExt;
2use reth_node_api::{BlockBody, PayloadKind};
3use reth_payload_builder::{PayloadBuilderHandle, PayloadId};
4use reth_payload_builder_primitives::Events;
5use reth_payload_primitives::{BuiltPayload, PayloadBuilderAttributes, PayloadTypes};
6use tokio_stream::wrappers::BroadcastStream;
7use tracing::debug;
8
9/// Helper for payload operations
10#[derive(derive_more::Debug)]
11pub struct PayloadTestContext<T: PayloadTypes> {
12    pub payload_event_stream: BroadcastStream<Events<T>>,
13    payload_builder: PayloadBuilderHandle<T>,
14    pub timestamp: u64,
15    #[debug(skip)]
16    attributes_generator: Box<dyn Fn(u64) -> T::PayloadBuilderAttributes + Send + Sync>,
17}
18
19impl<T: PayloadTypes> PayloadTestContext<T> {
20    /// Creates a new payload helper
21    pub async fn new(
22        payload_builder: PayloadBuilderHandle<T>,
23        attributes_generator: impl Fn(u64) -> T::PayloadBuilderAttributes + Send + Sync + 'static,
24    ) -> eyre::Result<Self> {
25        let payload_events = payload_builder.subscribe().await?;
26        let payload_event_stream = payload_events.into_stream();
27        // Cancun timestamp
28        Ok(Self {
29            payload_event_stream,
30            payload_builder,
31            timestamp: 1710338135,
32            attributes_generator: Box::new(attributes_generator),
33        })
34    }
35
36    /// Creates a new payload job from static attributes
37    pub async fn new_payload(&mut self) -> eyre::Result<T::PayloadBuilderAttributes> {
38        self.timestamp += 1;
39        let attributes = (self.attributes_generator)(self.timestamp);
40        self.payload_builder.send_new_payload(attributes.clone()).await.unwrap()?;
41        Ok(attributes)
42    }
43
44    /// Asserts that the next event is a payload attributes event
45    pub async fn expect_attr_event(
46        &mut self,
47        attrs: T::PayloadBuilderAttributes,
48    ) -> eyre::Result<()> {
49        let first_event = self.payload_event_stream.next().await.unwrap()?;
50        if let Events::Attributes(attr) = first_event {
51            assert_eq!(attrs.timestamp(), attr.timestamp());
52        } else {
53            panic!("Expect first event as payload attributes.")
54        }
55        Ok(())
56    }
57
58    /// Wait until the best built payload is ready
59    pub async fn wait_for_built_payload(&self, payload_id: PayloadId) {
60        loop {
61            let payload = self.payload_builder.best_payload(payload_id).await.unwrap().unwrap();
62            if payload.block().body().transactions().is_empty() {
63                tokio::time::sleep(std::time::Duration::from_millis(20)).await;
64                continue
65            }
66            // Resolve payload once its built
67            self.payload_builder
68                .resolve_kind(payload_id, PayloadKind::Earliest)
69                .await
70                .unwrap()
71                .unwrap();
72            break;
73        }
74    }
75
76    /// Expects the next event to be a built payload event or panics
77    pub async fn expect_built_payload(&mut self) -> eyre::Result<T::BuiltPayload> {
78        let second_event = self.payload_event_stream.next().await.unwrap()?;
79        if let Events::BuiltPayload(payload) = second_event {
80            Ok(payload)
81        } else {
82            panic!("Expect a built payload event.");
83        }
84    }
85}
reth_e2e_test_utils/payload.rs

reth_e2e_test_utils/
payload.rs
