third_party/rust/tokio-stream/tests/stream_merge.rs

   1 use tokio_stream::{self as stream, Stream, StreamExt};
   2 use tokio_test::task;
   3 use tokio_test::{assert_pending, assert_ready};
   4
   5 mod support {
   6     pub(crate) mod mpsc;
   7 }
   8
   9 use support::mpsc;
  10
  11 #[tokio::test]
  12 async fn merge_sync_streams() {
  13     let mut s = stream::iter(vec![0, 2, 4, 6]).merge(stream::iter(vec![1, 3, 5]));
  14
  15     for i in 0..7 {
  16         let rem = 7 - i;
  17         assert_eq!(s.size_hint(), (rem, Some(rem)));
  18         assert_eq!(Some(i), s.next().await);
  19     }
  20
  21     assert!(s.next().await.is_none());
  22 }
  23
  24 #[tokio::test]
  25 async fn merge_async_streams() {
  26     let (tx1, rx1) = mpsc::unbounded_channel_stream();
  27     let (tx2, rx2) = mpsc::unbounded_channel_stream();
  28
  29     let mut rx = task::spawn(rx1.merge(rx2));
  30
  31     assert_eq!(rx.size_hint(), (0, None));
  32
  33     assert_pending!(rx.poll_next());
  34
  35     tx1.send(1).unwrap();
  36
  37     assert!(rx.is_woken());
  38     assert_eq!(Some(1), assert_ready!(rx.poll_next()));
  39
  40     assert_pending!(rx.poll_next());
  41     tx2.send(2).unwrap();
  42
  43     assert!(rx.is_woken());
  44     assert_eq!(Some(2), assert_ready!(rx.poll_next()));
  45     assert_pending!(rx.poll_next());
  46
  47     drop(tx1);
  48     assert!(rx.is_woken());
  49     assert_pending!(rx.poll_next());
  50
  51     tx2.send(3).unwrap();
  52     assert!(rx.is_woken());
  53     assert_eq!(Some(3), assert_ready!(rx.poll_next()));
  54     assert_pending!(rx.poll_next());
  55
  56     drop(tx2);
  57     assert!(rx.is_woken());
  58     assert_eq!(None, assert_ready!(rx.poll_next()));
  59 }
  60
  61 #[test]
  62 fn size_overflow() {
  63     struct Monster;
  64
  65     impl tokio_stream::Stream for Monster {
  66         type Item = ();
  67         fn poll_next(
  68             self: std::pin::Pin<&mut Self>,
  69             _cx: &mut std::task::Context<'_>,
  70         ) -> std::task::Poll<Option<()>> {
  71             panic!()
  72         }
  73
  74         fn size_hint(&self) -> (usize, Option<usize>) {
  75             (usize::MAX, Some(usize::MAX))
  76         }
  77     }
  78
  79     let m1 = Monster;
  80     let m2 = Monster;
  81     let m = m1.merge(m2);
  82     assert_eq!(m.size_hint(), (usize::MAX, None));
  83 }