diff --git a/crates/optimism/flashblocks/src/ws/stream.rs b/crates/optimism/flashblocks/src/ws/stream.rs index e1f93fc7eab..faeda7c3564 100644 --- a/crates/optimism/flashblocks/src/ws/stream.rs +++ b/crates/optimism/flashblocks/src/ws/stream.rs @@ -181,7 +181,7 @@ mod tests { use crate::ExecutionPayloadBaseV1; use alloy_primitives::bytes::Bytes; use brotli::enc::BrotliEncoderParams; - use std::future; + use std::{future, iter}; use tokio_tungstenite::tungstenite::{Error, Utf8Bytes}; /// A `FakeConnector` creates [`FakeStream`]. @@ -238,6 +238,21 @@ mod tests { } } + /// Repeatedly fails to connect with the given error message. + #[derive(Clone)] + struct FailingConnector(String); + + impl WsConnect for FailingConnector { + type Stream = FakeStream; + + fn connect( + &mut self, + _ws_url: Url, + ) -> impl Future> + Send + Sync { + future::ready(Err(eyre!("{}", &self.0))) + } + } + fn to_json_message(block: &FlashBlock) -> Result { Ok(Message::Binary(Bytes::from(serde_json::to_vec(block).unwrap()))) } @@ -314,4 +329,19 @@ mod tests { assert_eq!(actual_messages, expected_messages); } + + #[tokio::test] + async fn test_connect_error_causes_retries() { + let tries = 3; + let error_msg = "test".to_owned(); + let messages = FailingConnector(error_msg.clone()); + let ws_url = "http://localhost".parse().unwrap(); + let stream = WsFlashBlockStream::with_connector(ws_url, messages); + + let actual_errors: Vec<_> = + stream.take(tries).map(Result::unwrap_err).map(|e| format!("{e}")).collect().await; + let expected_errors: Vec<_> = iter::repeat_n(error_msg, tries).collect(); + + assert_eq!(actual_errors, expected_errors); + } }