tokio/tokio-stream/tests/stream_fuse.rs
Lucio Franco 8efa62013b
Move stream items into tokio-stream (#3277)
This change removes all references to `Stream` from
within the `tokio` crate and moves them into a new
`tokio-stream` crate. Most types have had their
`impl Stream` removed as well in-favor of their
inherent methods.

Closes #2870
2020-12-15 20:24:38 -08:00

51 lines
1.3 KiB
Rust

use tokio_stream::{Stream, StreamExt};
use std::pin::Pin;
use std::task::{Context, Poll};
// a stream which alternates between Some and None
struct Alternate {
state: i32,
}
impl Stream for Alternate {
type Item = i32;
fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<i32>> {
let val = self.state;
self.state += 1;
// if it's even, Some(i32), else None
if val % 2 == 0 {
Poll::Ready(Some(val))
} else {
Poll::Ready(None)
}
}
}
#[tokio::test]
async fn basic_usage() {
let mut stream = Alternate { state: 0 };
// the stream goes back and forth
assert_eq!(stream.next().await, Some(0));
assert_eq!(stream.next().await, None);
assert_eq!(stream.next().await, Some(2));
assert_eq!(stream.next().await, None);
// however, once it is fused
let mut stream = stream.fuse();
assert_eq!(stream.size_hint(), (0, None));
assert_eq!(stream.next().await, Some(4));
assert_eq!(stream.size_hint(), (0, None));
assert_eq!(stream.next().await, None);
// it will always return `None` after the first time.
assert_eq!(stream.size_hint(), (0, Some(0)));
assert_eq!(stream.next().await, None);
assert_eq!(stream.size_hint(), (0, Some(0)));
}