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]
merge_sync_streams()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]
merge_async_streams()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]
size_overflow()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 }
84