Fix flaky tests (zio#914)

7mind · May 30, 2019 · 6ef36b1 · 6ef36b1
1 parent 7342338
commit 6ef36b1
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 19 deletions.
diff --git a/core/jvm/src/test/scala/scalaz/zio/RTSSpec.scala b/core/jvm/src/test/scala/scalaz/zio/RTSSpec.scala
@@ -968,7 +968,7 @@ class RTSSpec(implicit ee: ExecutionEnv) extends TestRuntime {
     } yield result must_=== 10
   }
 
-  def testInterruptStatusIsHeritable = unsafeRun {
+  def testInterruptStatusIsHeritable = nonFlaky {
     for {
       latch <- Promise.make[Nothing, Unit]
       ref   <- Ref.make(InterruptStatus.interruptible)
@@ -977,16 +977,15 @@ class RTSSpec(implicit ee: ExecutionEnv) extends TestRuntime {
     } yield v must_=== InterruptStatus.uninterruptible
   }
 
-  def testExecutorIsHeritable = unsafeRun {
-    for {
+  def testExecutorIsHeritable =
+    nonFlaky(for {
       ref  <- Ref.make(Option.empty[internal.Executor])
       exec = internal.Executor.fromExecutionContext(100)(scala.concurrent.ExecutionContext.Implicits.global)
-      _    <- IO.descriptor.map(_.executor).flatMap(e => ref.set(Some(e))).fork.lock(exec)
+      _    <- withLatch(release => IO.descriptor.map(_.executor).flatMap(e => ref.set(Some(e)) *> release).fork.lock(exec))
       v    <- ref.get
-    } yield v must_=== Some(exec)
-  }
+    } yield v must_=== Some(exec))
 
-  def testSupervisionIsHeritable = unsafeRun {
+  def testSupervisionIsHeritable = nonFlaky {
     for {
       latch <- Promise.make[Nothing, Unit]
       ref   <- Ref.make(SuperviseStatus.unsupervised)
@@ -996,18 +995,15 @@ class RTSSpec(implicit ee: ExecutionEnv) extends TestRuntime {
   }
 
   def testSupervisingInheritance = {
-    def forkAwaitStart[A](io: UIO[A]) =
-      for {
-        latch <- Promise.make[Nothing, Unit]
-        _     <- (latch.succeed(()) *> io).fork
-        _     <- latch.await
-      } yield ()
+    def forkAwaitStart[A](io: UIO[A], refs: Ref[List[Fiber[_, _]]]): UIO[Fiber[Nothing, A]] =
+      withLatch(release => (release *> io).fork.tap(f => refs.update(f :: _)))
 
-    unsafeRun(
+    nonFlaky(
       (for {
-        _    <- forkAwaitStart(forkAwaitStart(forkAwaitStart(IO.succeed(()))))
+        ref  <- Ref.make[List[Fiber[_, _]]](Nil) // To make strong ref
+        _    <- forkAwaitStart(forkAwaitStart(forkAwaitStart(IO.succeed(()), ref), ref), ref)
         fibs <- ZIO.children
-      } yield fibs must have size 3).supervised
+      } yield fibs must have size 1).supervised
     )
   }
 
@@ -1435,4 +1431,10 @@ class RTSSpec(implicit ee: ExecutionEnv) extends TestRuntime {
     unsafeRun(
       IO.mergeAll(List.empty[UIO[Int]])(0)(_ + _)
     ) must_=== 0
+
+  def nonFlaky(v: => ZIO[Environment, Any, org.specs2.matcher.MatchResult[_]]): org.specs2.matcher.MatchResult[_] =
+    (1 to 50).foldLeft[org.specs2.matcher.MatchResult[_]](true must_=== true) {
+      case (acc, _) =>
+        acc and unsafeRun(v)
+    }
 }
diff --git a/core/shared/src/main/scala/scalaz/zio/internal/FiberContext.scala b/core/shared/src/main/scala/scalaz/zio/internal/FiberContext.scala
@@ -378,12 +378,17 @@ private[zio] final class FiberContext[E, A](
   final def fork[E, A](zio: IO[E, A]): FiberContext[E, A] = {
     val childFiberRefLocals: FiberRefLocals = platform.newWeakHashMap()
     childFiberRefLocals.putAll(fiberRefLocals)
+
+    val childSupervised = supervised.peek() match {
+      case SuperviseStatus.Unsupervised  => SuperviseStatus.Unsupervised
+      case SuperviseStatus.Supervised(_) => SuperviseStatus.Supervised(newWeakSet[Fiber[_, _]])
+    }
     val context = new FiberContext[E, A](
       platform,
       environments.peek(),
       executors.peek(),
       InterruptStatus.fromBoolean(interruptStatus.peekOrElse(true)),
-      supervised.peek(),
+      childSupervised,
       childFiberRefLocals
     )
 
@@ -423,11 +428,12 @@ private[zio] final class FiberContext[E, A](
       }
   }
 
+  private[this] final def newWeakSet[A]: Set[A] =
+    Collections.newSetFromMap[A](platform.newWeakHashMap[A, java.lang.Boolean]())
+
   private[this] final def changeSupervision(status: scalaz.zio.SuperviseStatus): IO[E, Unit] = ZIO.effectTotal {
     status match {
       case scalaz.zio.SuperviseStatus.Supervised =>
-        def newWeakSet[A]: Set[A] = Collections.newSetFromMap[A](platform.newWeakHashMap[A, java.lang.Boolean]())
-
         val set = newWeakSet[Fiber[_, _]]
 
         supervised.push(SuperviseStatus.Supervised(set))