stream: stop polling the underlying stream once map_while yields None (#8233)

tokio-stream/tests/stream_map_while.rs

@@ -0,0 +1,22 @@
+use tokio_stream::StreamExt;
+
+#[tokio::test]
+async fn map_while_yields_until_closure_returns_none() {
+    let mut stream =
+        tokio_stream::iter(1..=10).map_while(|x| if x < 4 { Some(x + 3) } else { None });
+    assert_eq!(stream.next().await, Some(4));
+    assert_eq!(stream.next().await, Some(5));
+    assert_eq!(stream.next().await, Some(6));
+    assert_eq!(stream.next().await, None);
+}
+
+#[tokio::test]
+async fn map_while_does_not_poll_after_closure_returns_none() {
+    // Once the closure returns `None`, the underlying stream must not be polled
+    // again, so the trailing `2` is never yielded.
+    let mut stream =
+        tokio_stream::iter(vec![1, 5, 2]).map_while(|x| if x < 3 { Some(x) } else { None });
+    assert_eq!(stream.next().await, Some(1));
+    assert_eq!(stream.next().await, None);
+    assert_eq!(stream.next().await, None);
+}