From 0b39306b747734997cfecca76a7ea912fae624ba Mon Sep 17 00:00:00 2001
From: Yoshua Wuyts <yoshuawuyts@gmail.com>
Date: Thu, 26 Sep 2019 17:24:24 +0200
Subject: [PATCH] fix barrier tests

Signed-off-by: Yoshua Wuyts <yoshuawuyts@gmail.com>
---
 src/sync/barrier.rs | 55 +++++++++++++++++++++++++++++++++++++++++++++
 tests/barrier.rs    | 52 ------------------------------------------
 2 files changed, 55 insertions(+), 52 deletions(-)
 delete mode 100644 tests/barrier.rs

diff --git a/src/sync/barrier.rs b/src/sync/barrier.rs
index d8981a0..44c94b8 100644
--- a/src/sync/barrier.rs
+++ b/src/sync/barrier.rs
@@ -174,3 +174,58 @@ impl BarrierWaitResult {
         self.0
     }
 }
+
+#[cfg(test)]
+mod test {
+    use futures_channel::mpsc::unbounded;
+    use futures_util::sink::SinkExt;
+    use futures_util::stream::StreamExt;
+
+    use crate::sync::{Arc, Barrier};
+    use crate::task;
+
+    #[test]
+    fn test_barrier() {
+        // NOTE(dignifiedquire): Based on the test in std, I was seeing some
+        // race conditions, so running it in a loop to make sure things are
+        // solid.
+
+        for _ in 0..1_000 {
+            task::block_on(async move {
+                const N: usize = 10;
+
+                let barrier = Arc::new(Barrier::new(N));
+                let (tx, mut rx) = unbounded();
+
+                for _ in 0..N - 1 {
+                    let c = barrier.clone();
+                    let mut tx = tx.clone();
+                    task::spawn(async move {
+                        let res = c.wait().await;
+
+                        tx.send(res.is_leader()).await.unwrap();
+                    });
+                }
+
+                // At this point, all spawned threads should be blocked,
+                // so we shouldn't get anything from the port
+                let res = rx.try_next();
+                assert!(match res {
+                    Err(_err) => true,
+                    _ => false,
+                });
+
+                let mut leader_found = barrier.wait().await.is_leader();
+
+                // Now, the barrier is cleared and we should get data.
+                for _ in 0..N - 1 {
+                    if rx.next().await.unwrap() {
+                        assert!(!leader_found);
+                        leader_found = true;
+                    }
+                }
+                assert!(leader_found);
+            });
+        }
+    }
+}
diff --git a/tests/barrier.rs b/tests/barrier.rs
deleted file mode 100644
index 3284944..0000000
--- a/tests/barrier.rs
+++ /dev/null
@@ -1,52 +0,0 @@
-use std::sync::Arc;
-
-use futures_channel::mpsc::unbounded;
-use futures_util::sink::SinkExt;
-use futures_util::stream::StreamExt;
-
-use async_std::sync::Barrier;
-use async_std::task;
-
-#[test]
-fn test_barrier() {
-    // Based on the test in std, I was seeing some race conditions, so running it in a loop to make sure
-    // things are solid.
-
-    for _ in 0..1_000 {
-        task::block_on(async move {
-            const N: usize = 10;
-
-            let barrier = Arc::new(Barrier::new(N));
-            let (tx, mut rx) = unbounded();
-
-            for _ in 0..N - 1 {
-                let c = barrier.clone();
-                let mut tx = tx.clone();
-                task::spawn(async move {
-                    let res = c.wait().await;
-
-                    tx.send(res.is_leader()).await.unwrap();
-                });
-            }
-
-            // At this point, all spawned threads should be blocked,
-            // so we shouldn't get anything from the port
-            let res = rx.try_next();
-            assert!(match res {
-                Err(_err) => true,
-                _ => false,
-            });
-
-            let mut leader_found = barrier.wait().await.is_leader();
-
-            // Now, the barrier is cleared and we should get data.
-            for _ in 0..N - 1 {
-                if rx.next().await.unwrap() {
-                    assert!(!leader_found);
-                    leader_found = true;
-                }
-            }
-            assert!(leader_found);
-        });
-    }
-}