1 use tokio_stream::{Stream, StreamExt};
2
3 use std::pin::Pin;
4 use std::task::{Context, Poll};
5
6 // a stream which alternates between Some and None
7 struct Alternate {
8 state: i32,
9 }
10
11 impl Stream for Alternate {
12 type Item = i32;
13
poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<i32>>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]
basic_usage()28 async 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