1 files changed, 330 insertions, 0 deletions
diff --git a/third_party/rust/tokio-stream/tests/stream_stream_map.rs b/third_party/rust/tokio-stream/tests/stream_stream_map.rs
new file mode 100644
index 0000000000..b6b87e9d0a
--- /dev/null
+++ b/third_party/rust/tokio-stream/tests/stream_stream_map.rs
@@ -0,0 +1,330 @@
+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)
+}