fix(optimism): Reconnect if ws stream ends in WsFlashBlockStream

Author: RomanHodulak

crates/optimism/flashblocks/src/ws/stream.rs

@@ -81,7 +81,7 @@ where
                 match ready!(this.connect.poll_unpin(cx)) {
                     Ok((sink, stream)) => this.stream(sink, stream),
                     Err(err) => {
-                        this.state = State::Initial;
+                        this.connect();
 
                         return Poll::Ready(Some(Err(err)));
                     }
@@ -104,7 +104,9 @@ where
                     .expect("Stream state should be unreachable without stream")
                     .poll_next_unpin(cx))
                 else {
-                    return Poll::Ready(None);
+                    this.connect();
+
+                    break;
                 };
 
                 match msg {
@@ -432,18 +434,43 @@ mod tests {
         let ws_url = "http://localhost".parse().unwrap();
         let stream = WsFlashBlockStream::with_connector(ws_url, messages);
 
-        let actual_messages: Vec<_> = stream.map(Result::unwrap).collect().await;
+        let actual_messages: Vec<_> = stream.take(1).map(Result::unwrap).collect().await;
         let expected_messages = flashblocks.to_vec();
 
         assert_eq!(actual_messages, expected_messages);
     }
 
     #[tokio::test]
     async fn test_stream_ignores_non_binary_message() {
-        let messages = FakeConnector::from([Ok(Message::Text(Utf8Bytes::from("test")))]);
+        let flashblock = FlashBlock {
+            payload_id: Default::default(),
+            index: 0,
+            base: Some(ExecutionPayloadBaseV1 {
+                parent_beacon_block_root: Default::default(),
+                parent_hash: Default::default(),
+                fee_recipient: Default::default(),
+                prev_randao: Default::default(),
+                block_number: 0,
+                gas_limit: 0,
+                timestamp: 0,
+                extra_data: Default::default(),
+                base_fee_per_gas: Default::default(),
+            }),
+            diff: Default::default(),
+            metadata: Default::default(),
+        };
+        let messages = FakeConnector::from([
+            Ok(Message::Text(Utf8Bytes::from("test"))),
+            to_json_message(&flashblock),
+        ]);
         let ws_url = "http://localhost".parse().unwrap();
         let mut stream = WsFlashBlockStream::with_connector(ws_url, messages);
-        assert!(stream.next().await.is_none());
+
+        let expected_message = flashblock;
+        let actual_message =
+            stream.next().await.expect("Binary message should not be ignored").unwrap();
+
+        assert_eq!(actual_message, expected_message)
     }
 
     #[tokio::test]
@@ -453,7 +480,7 @@ mod tests {
         let stream = WsFlashBlockStream::with_connector(ws_url, messages);
 
         let actual_messages: Vec<_> =
-            stream.map(Result::unwrap_err).map(|e| format!("{e}")).collect().await;
+            stream.take(1).map(Result::unwrap_err).map(|e| format!("{e}")).collect().await;
         let expected_messages = vec!["Attack attempt detected".to_owned()];
 
         assert_eq!(actual_messages, expected_messages);