reth_e2e_test_utils/
payload.rs1use 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#[derive(derive_more::Debug)]
10pub struct PayloadTestContext<T: PayloadTypes> {
11 pub payload_event_stream: BroadcastStream<Events<T>>,
12 payload_builder: PayloadBuilderHandle<T>,
13 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 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 Ok(Self {
35 payload_event_stream,
36 payload_builder,
37 timestamp: 1710338135,
38 attributes_generator: Box::new(attributes_generator),
39 })
40 }
41
42 pub fn next_attributes(&mut self) -> T::PayloadAttributes {
44 self.timestamp += 1;
45 (self.attributes_generator)(self.timestamp)
46 }
47
48 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 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 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}