1use tokio_stream::{Stream, StreamExt};
2
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6// a stream which alternates between Some and None
7struct Alternate {
8 state: i32,
9}
10
11impl Stream for Alternate {
12 type Item = i32;
13
14 fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<i32>> {
15 let val = self.state;
16 self.state += 1;
17
18 // if it's even, Some(i32), else None
19 if val % 2 == 0 {
20 Poll::Ready(Some(val))
21 } else {
22 Poll::Ready(None)
23 }
24 }
25}
26
27#[tokio::test]
28async fn basic_usage() {
29 let mut stream = Alternate { state: 0 };
30
31 // the stream goes back and forth
32 assert_eq!(stream.next().await, Some(0));
33 assert_eq!(stream.next().await, None);
34 assert_eq!(stream.next().await, Some(2));
35 assert_eq!(stream.next().await, None);
36
37 // however, once it is fused
38 let mut stream = stream.fuse();
39
40 assert_eq!(stream.size_hint(), (0, None));
41 assert_eq!(stream.next().await, Some(4));
42
43 assert_eq!(stream.size_hint(), (0, None));
44 assert_eq!(stream.next().await, None);
45
46 // it will always return `None` after the first time.
47 assert_eq!(stream.size_hint(), (0, Some(0)));
48 assert_eq!(stream.next().await, None);
49 assert_eq!(stream.size_hint(), (0, Some(0)));
50}
51