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;
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,
14 #[debug(skip)]
15 attributes_generator: Box<dyn Fn(u64) -> T::PayloadBuilderAttributes + Send + Sync>,
16}
17
18impl<T: PayloadTypes> PayloadTestContext<T> {
19 pub async fn new(
21 payload_builder: PayloadBuilderHandle<T>,
22 attributes_generator: impl Fn(u64) -> T::PayloadBuilderAttributes + Send + Sync + 'static,
23 ) -> eyre::Result<Self> {
24 let payload_events = payload_builder.subscribe().await?;
25 let payload_event_stream = payload_events.into_stream();
26 Ok(Self {
28 payload_event_stream,
29 payload_builder,
30 timestamp: 1710338135,
31 attributes_generator: Box::new(attributes_generator),
32 })
33 }
34
35 pub async fn new_payload(&mut self) -> eyre::Result<T::PayloadBuilderAttributes> {
37 self.timestamp += 1;
38 let attributes = (self.attributes_generator)(self.timestamp);
39 self.payload_builder.send_new_payload(attributes.clone()).await.unwrap()?;
40 Ok(attributes)
41 }
42
43 pub async fn expect_attr_event(
45 &mut self,
46 attrs: T::PayloadBuilderAttributes,
47 ) -> eyre::Result<()> {
48 let first_event = self.payload_event_stream.next().await.unwrap()?;
49 if let Events::Attributes(attr) = first_event {
50 assert_eq!(attrs.timestamp(), attr.timestamp());
51 } else {
52 panic!("Expect first event as payload attributes.")
53 }
54 Ok(())
55 }
56
57 pub async fn wait_for_built_payload(&self, payload_id: PayloadId) {
59 loop {
60 let payload = self.payload_builder.best_payload(payload_id).await.unwrap().unwrap();
61 if payload.block().body().transactions().is_empty() {
62 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
63 continue
64 }
65 self.payload_builder
67 .resolve_kind(payload_id, PayloadKind::Earliest)
68 .await
69 .unwrap()
70 .unwrap();
71 break;
72 }
73 }
74
75 pub async fn expect_built_payload(&mut self) -> eyre::Result<T::BuiltPayload> {
77 let second_event = self.payload_event_stream.next().await.unwrap()?;
78 if let Events::BuiltPayload(payload) = second_event {
79 Ok(payload)
80 } else {
81 panic!("Expect a built payload event.");
82 }
83 }
84}