Source code
Revision control
Copy as Markdown
Other Tools
use tokio_stream::{self as stream, pending, Stream, StreamExt, StreamMap};
use tokio_test::{assert_ok, assert_pending, assert_ready, task};
mod support {
pub(crate) mod mpsc;
}
use support::mpsc;
use std::pin::Pin;
macro_rules! assert_ready_some {
($($t:tt)*) => {
match assert_ready!($($t)*) {
Some(v) => v,
None => panic!("expected `Some`, got `None`"),
}
};
}
macro_rules! assert_ready_none {
($($t:tt)*) => {
match assert_ready!($($t)*) {
None => {}
Some(v) => panic!("expected `None`, got `Some({:?})`", v),
}
};
}
#[tokio::test]
async fn empty() {
let mut map = StreamMap::<&str, stream::Pending<()>>::new();
assert_eq!(map.len(), 0);
assert!(map.is_empty());
assert!(map.next().await.is_none());
assert!(map.next().await.is_none());
assert!(map.remove("foo").is_none());
}
#[tokio::test]
async fn single_entry() {
let mut map = task::spawn(StreamMap::new());
let (tx, rx) = mpsc::unbounded_channel_stream();
let rx = Box::pin(rx);
assert_ready_none!(map.poll_next());
assert!(map.insert("foo", rx).is_none());
assert!(map.contains_key("foo"));
assert!(!map.contains_key("bar"));
assert_eq!(map.len(), 1);
assert!(!map.is_empty());
assert_pending!(map.poll_next());
assert_ok!(tx.send(1));
assert!(map.is_woken());
let (k, v) = assert_ready_some!(map.poll_next());
assert_eq!(k, "foo");
assert_eq!(v, 1);
assert_pending!(map.poll_next());
assert_ok!(tx.send(2));
assert!(map.is_woken());
let (k, v) = assert_ready_some!(map.poll_next());
assert_eq!(k, "foo");
assert_eq!(v, 2);
assert_pending!(map.poll_next());
drop(tx);
assert!(map.is_woken());
assert_ready_none!(map.poll_next());
}
#[tokio::test]
async fn multiple_entries() {
let mut map = task::spawn(StreamMap::new());
let (tx1, rx1) = mpsc::unbounded_channel_stream();
let (tx2, rx2) = mpsc::unbounded_channel_stream();
let rx1 = Box::pin(rx1);
let rx2 = Box::pin(rx2);
map.insert("foo", rx1);
map.insert("bar", rx2);
assert_pending!(map.poll_next());
assert_ok!(tx1.send(1));
assert!(map.is_woken());
let (k, v) = assert_ready_some!(map.poll_next());
assert_eq!(k, "foo");
assert_eq!(v, 1);
assert_pending!(map.poll_next());
assert_ok!(tx2.send(2));
assert!(map.is_woken());
let (k, v) = assert_ready_some!(map.poll_next());
assert_eq!(k, "bar");
assert_eq!(v, 2);
assert_pending!(map.poll_next());
assert_ok!(tx1.send(3));
assert_ok!(tx2.send(4));
assert!(map.is_woken());
// Given the randomization, there is no guarantee what order the values will
// be received in.
let mut v = (0..2)
.map(|_| assert_ready_some!(map.poll_next()))
.collect::<Vec<_>>();
assert_pending!(map.poll_next());
v.sort_unstable();
assert_eq!(v[0].0, "bar");
assert_eq!(v[0].1, 4);
assert_eq!(v[1].0, "foo");
assert_eq!(v[1].1, 3);
drop(tx1);
assert!(map.is_woken());
assert_pending!(map.poll_next());
drop(tx2);
assert_ready_none!(map.poll_next());
}
#[tokio::test]
async fn insert_remove() {
let mut map = task::spawn(StreamMap::new());
let (tx, rx) = mpsc::unbounded_channel_stream();
let rx = Box::pin(rx);
assert_ready_none!(map.poll_next());
assert!(map.insert("foo", rx).is_none());
let rx = map.remove("foo").unwrap();
assert_ok!(tx.send(1));
assert!(!map.is_woken());
assert_ready_none!(map.poll_next());
assert!(map.insert("bar", rx).is_none());
let v = assert_ready_some!(map.poll_next());
assert_eq!(v.0, "bar");
assert_eq!(v.1, 1);
assert!(map.remove("bar").is_some());
assert_ready_none!(map.poll_next());
assert!(map.is_empty());
assert_eq!(0, map.len());
}
#[tokio::test]
async fn replace() {
let mut map = task::spawn(StreamMap::new());
let (tx1, rx1) = mpsc::unbounded_channel_stream();
let (tx2, rx2) = mpsc::unbounded_channel_stream();
let rx1 = Box::pin(rx1);
let rx2 = Box::pin(rx2);
assert!(map.insert("foo", rx1).is_none());
assert_pending!(map.poll_next());
let _rx1 = map.insert("foo", rx2).unwrap();
assert_pending!(map.poll_next());
tx1.send(1).unwrap();
assert_pending!(map.poll_next());
tx2.send(2).unwrap();
assert!(map.is_woken());
let v = assert_ready_some!(map.poll_next());
assert_eq!(v.0, "foo");
assert_eq!(v.1, 2);
}
#[test]
fn size_hint_with_upper() {
let mut map = StreamMap::new();
map.insert("a", stream::iter(vec![1]));
map.insert("b", stream::iter(vec![1, 2]));
map.insert("c", stream::iter(vec![1, 2, 3]));
assert_eq!(3, map.len());
assert!(!map.is_empty());
let size_hint = map.size_hint();
assert_eq!(size_hint, (6, Some(6)));
}
#[test]
fn size_hint_without_upper() {
let mut map = StreamMap::new();
map.insert("a", pin_box(stream::iter(vec![1])));
map.insert("b", pin_box(stream::iter(vec![1, 2])));
map.insert("c", pin_box(pending()));
let size_hint = map.size_hint();
assert_eq!(size_hint, (3, None));
}
#[test]
fn new_capacity_zero() {
let map = StreamMap::<&str, stream::Pending<()>>::new();
assert_eq!(0, map.capacity());
assert!(map.keys().next().is_none());
}
#[test]
fn with_capacity() {
let map = StreamMap::<&str, stream::Pending<()>>::with_capacity(10);
assert!(10 <= map.capacity());
assert!(map.keys().next().is_none());
}
#[test]
fn iter_keys() {
let mut map = StreamMap::new();
map.insert("a", pending::<i32>());
map.insert("b", pending());
map.insert("c", pending());
let mut keys = map.keys().collect::<Vec<_>>();
keys.sort_unstable();
assert_eq!(&keys[..], &[&"a", &"b", &"c"]);
}
#[test]
fn iter_values() {
let mut map = StreamMap::new();
map.insert("a", stream::iter(vec![1]));
map.insert("b", stream::iter(vec![1, 2]));
map.insert("c", stream::iter(vec![1, 2, 3]));
let mut size_hints = map.values().map(|s| s.size_hint().0).collect::<Vec<_>>();
size_hints.sort_unstable();
assert_eq!(&size_hints[..], &[1, 2, 3]);
}
#[test]
fn iter_values_mut() {
let mut map = StreamMap::new();
map.insert("a", stream::iter(vec![1]));
map.insert("b", stream::iter(vec![1, 2]));
map.insert("c", stream::iter(vec![1, 2, 3]));
let mut size_hints = map
.values_mut()
.map(|s: &mut _| s.size_hint().0)
.collect::<Vec<_>>();
size_hints.sort_unstable();
assert_eq!(&size_hints[..], &[1, 2, 3]);
}
#[test]
fn clear() {
let mut map = task::spawn(StreamMap::new());
map.insert("a", stream::iter(vec![1]));
map.insert("b", stream::iter(vec![1, 2]));
map.insert("c", stream::iter(vec![1, 2, 3]));
assert_ready_some!(map.poll_next());
map.clear();
assert_ready_none!(map.poll_next());
assert!(map.is_empty());
}
#[test]
fn contains_key_borrow() {
let mut map = StreamMap::new();
map.insert("foo".to_string(), pending::<()>());
assert!(map.contains_key("foo"));
}
#[test]
fn one_ready_many_none() {
// Run a few times because of randomness
for _ in 0..100 {
let mut map = task::spawn(StreamMap::new());
map.insert(0, pin_box(stream::empty()));
map.insert(1, pin_box(stream::empty()));
map.insert(2, pin_box(stream::once("hello")));
map.insert(3, pin_box(stream::pending()));
let v = assert_ready_some!(map.poll_next());
assert_eq!(v, (2, "hello"));
}
}
fn pin_box<T: Stream<Item = U> + 'static, U>(s: T) -> Pin<Box<dyn Stream<Item = U>>> {
Box::pin(s)
}