From 5dda7fd95e40e77a9159dbe753b5a35d8f7f2c42 Mon Sep 17 00:00:00 2001
From: Mark Lodato <marklodato0@gmail.com>
Date: Sun, 30 Jun 2024 20:37:48 -0400
Subject: [PATCH] Allow tests to be run with Miri

## This Commit

Leverages [this comment][0] to allow tests to be run with Miri.

## Why?

The regular macro enables a runtime that cannot be run with Miri yet and
we'd like to catch any undefined behavior we can.

[0]: https://github.com/rust-lang/miri/issues/602#issuecomment-884019764
---
 src/lib.rs | 123 ++++++++++++++++++++++++++++++-----------------------
 1 file changed, 69 insertions(+), 54 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 5b43577..387a975 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,7 +3,6 @@
 //    - Consolidate structs with const generics? (pin_project doesn't handle this)
 // 3. Better tests
 //    - Loom
-//    - Miri
 //    - Quickcheck/proptest
 //    - tokio::test to skip time
 // 4. Performance optimizations
@@ -136,69 +135,85 @@ mod tests {
     use std::time::Duration;
 
     use async_stream::stream;
-    use futures::{FutureExt, StreamExt};
+    use futures::{Future, FutureExt, StreamExt};
     use itertools::Itertools;
 
-    #[tokio::test]
-    async fn outerleaves() {
-        let stream = futures::stream::iter([0, 1, 2, 3, 4, 5, 6]);
-        let (mut evens, mut odds) = stream.outerleave();
+    // See https://github.com/rust-lang/miri/issues/602#issuecomment-884019764
+    fn miri_test(f: impl Future<Output = ()>) {
+        let rt = tokio::runtime::Builder::new_current_thread()
+            .enable_time()
+            .build()
+            .unwrap();
 
-        assert_eq!(evens.next().await, Some(0));
-        assert_eq!(odds.next().await, Some(1));
-        assert_eq!(evens.next().await, Some(2));
-
-        assert_eq!(evens.next().now_or_never(), None);
-
-        let jh = tokio::spawn(async move {
-            assert_eq!(evens.next().await, Some(4));
-        });
-        assert_eq!(odds.next().await, Some(3));
-        jh.await.unwrap();
+        rt.block_on(f)
     }
 
-    #[tokio::test]
-    async fn handles_stream_not_ready() {
-        let stream = futures::stream::iter([0, 1]);
-        let (mut evens, mut odds) = stream.outerleave();
+    #[test]
+    fn outerleaves() {
+        miri_test(async move {
+            let stream = futures::stream::iter([0, 1, 2, 3, 4, 5, 6]);
+            let (mut evens, mut odds) = stream.outerleave();
 
-        let jh = tokio::spawn(async move { odds.next().await });
+            assert_eq!(evens.next().await, Some(0));
+            assert_eq!(odds.next().await, Some(1));
+            assert_eq!(evens.next().await, Some(2));
 
-        tokio::time::sleep(Duration::from_millis(10)).await;
-        assert_eq!(evens.next().await, Some(0));
+            assert_eq!(evens.next().now_or_never(), None);
 
-        assert_eq!(
-            tokio::time::timeout(Duration::from_millis(10), jh)
-                .await
-                .unwrap()
-                .unwrap(),
-            Some(1)
-        );
+            let jh = tokio::spawn(async move {
+                assert_eq!(evens.next().await, Some(4));
+            });
+            assert_eq!(odds.next().await, Some(3));
+            jh.await.unwrap();
+        });
     }
 
-    #[tokio::test]
+    #[test]
+    fn handles_stream_not_ready() {
+        miri_test(async move {
+            let stream = futures::stream::iter([0, 1]);
+            let (mut evens, mut odds) = stream.outerleave();
+
+            let jh = tokio::spawn(async move { odds.next().await });
+
+            tokio::time::sleep(Duration::from_millis(10)).await;
+            assert_eq!(evens.next().await, Some(0));
+
+            assert_eq!(
+                tokio::time::timeout(Duration::from_millis(10), jh)
+                    .await
+                    .unwrap()
+                    .unwrap(),
+                Some(1)
+            );
+        })
+    }
+
+    #[test]
     #[ignore = "takes a long time"]
-    async fn handles_different_yielding_patterns() {
-        let sleeps: Vec<_> = (0..6).map(|n| Duration::from_millis(n * 10)).collect();
-        let sleeps = sleeps.iter().copied().permutations(sleeps.len()).take(300);
-        for sleeps in sleeps {
-            let stream = stream! {
-                for (i, sleep) in (0..6).zip(sleeps) {
-                    tokio::time::sleep(sleep).await;
-                    yield Ok(i);
-                }
-            };
-
-            let (evens, odds) = stream.outerleave();
-            let (tx, rx) = futures::channel::mpsc::unbounded();
-            let even_fut = evens.forward(tx.clone());
-            let odd_fut = odds.forward(tx);
-            let (even_result, odd_result) = futures::join!(even_fut, odd_fut);
-            even_result.unwrap();
-            odd_result.unwrap();
-
-            let received: Vec<_> = rx.collect().await;
-            assert_eq!(received, [0, 1, 2, 3, 4, 5]);
-        }
+    fn handles_different_yielding_patterns() {
+        miri_test(async move {
+            let sleeps: Vec<_> = (0..6).map(|n| Duration::from_millis(n * 10)).collect();
+            let sleeps = sleeps.iter().copied().permutations(sleeps.len()).take(300);
+            for sleeps in sleeps {
+                let stream = stream! {
+                    for (i, sleep) in (0..6).zip(sleeps) {
+                        tokio::time::sleep(sleep).await;
+                        yield Ok(i);
+                    }
+                };
+
+                let (evens, odds) = stream.outerleave();
+                let (tx, rx) = futures::channel::mpsc::unbounded();
+                let even_fut = evens.forward(tx.clone());
+                let odd_fut = odds.forward(tx);
+                let (even_result, odd_result) = futures::join!(even_fut, odd_fut);
+                even_result.unwrap();
+                odd_result.unwrap();
+
+                let received: Vec<_> = rx.collect().await;
+                assert_eq!(received, [0, 1, 2, 3, 4, 5]);
+            }
+        });
     }
 }