test(optimism): Cover the case of repeatedly failing to connect to websocket in WsFlashBlockStream

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<Output = eyre::Result<Self::Stream>> + Send + Sync {
+            future::ready(Err(eyre!("{}", &self.0)))
+        }
+    }
+
     fn to_json_message(block: &FlashBlock) -> Result<Message, Error> {
         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);
+    }
 }