Fix race condition in groupedWeighted* tests (#171)

softwaremill · Jul 3, 2024 · a32e6c6 · a32e6c6
1 parent 2121658
commit a32e6c6
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/core/src/test/scala/ox/channels/SourceOpsGroupedTest.scala b/core/src/test/scala/ox/channels/SourceOpsGroupedTest.scala
@@ -37,9 +37,8 @@ class SourceOpsGroupedTest extends AnyFlatSpec with Matchers {
   }
 
   it should "return failed source when cost function throws exception" in supervised {
-    val source = Source.fromValues(1, 2, 3, 0, 4, 5, 6, 7).groupedWeighted(150)(n => 100 / n)
-    source.receiveOrClosed() shouldBe Seq(1, 2)
-    source.receiveOrClosed() should matchPattern {
+    val result = Source.fromValues(1, 2, 3, 0, 4, 5, 6, 7).groupedWeighted(150)(n => 100 / n).drainOrError()
+    result should matchPattern {
       case ChannelClosed.Error(reason) if reason.isInstanceOf[ArithmeticException] =>
     }
   }
@@ -154,9 +153,8 @@ class SourceOpsGroupedTest extends AnyFlatSpec with Matchers {
   }
 
   it should "return failed source when cost function throws exception" in supervised {
-    val source = Source.fromValues(1, 2, 3, 0, 4, 5, 6, 7).groupedWeightedWithin(150, 100.millis)(n => 100 / n)
-    source.receiveOrClosed() shouldBe Seq(1, 2)
-    source.receiveOrClosed() should matchPattern {
+    val result = Source.fromValues(1, 2, 3, 0, 4, 5, 6, 7).groupedWeightedWithin(150, 100.millis)(n => 100 / n).drainOrError()
+    result should matchPattern {
       case ChannelClosed.Error(reason) if reason.isInstanceOf[ArithmeticException] =>
     }
   }