Fiber resource

sky-uk · Jun 29, 2023 · 55e77c4 · 55e77c4
1 parent 67de72a
commit 55e77c4
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 16 deletions.
diff --git a/src/test/scala/integration/TopicLoaderIntSpec.scala b/src/test/scala/integration/TopicLoaderIntSpec.scala
@@ -189,25 +189,25 @@ class TopicLoaderIntSpec extends KafkaSpecBase[IO] {
 
       val (preLoad, postLoad) = records(1 to 15).splitAt(10)
 
-      val loadRef: IO[Ref[IO, Boolean]]                = Ref.of(false)
-      val topicRef: IO[Ref[IO, Seq[(String, String)]]] = Ref.empty
-
       for {
-        loadState  <- loadRef
-        topicState <- topicRef
+        loadState  <- Ref.of[IO, Boolean](false)
+        topicState <- Ref.empty[IO, Seq[(String, String)]]
         _          <- createCustomTopics(NonEmptyList.one(testTopic1))
         _          <- publishStringMessages(testTopic1, preLoad)
-        fiber      <- loadAndRunLoader(NonEmptyList.one(testTopic1))(_ => loadState.set(true))
-                        .map(recordToTuple)
-                        .evalTap(r => topicState.getAndUpdate(_ :+ r))
-                        .compile
-                        .drain
-                        .start
-        _          <- eventually(topicState.get.asserting(_ should contain theSameElementsAs preLoad))
-        _          <- loadState.get.asserting(_ shouldBe true)
-        _          <- publishStringMessages(testTopic1, postLoad)
-        assertion  <- eventually(topicState.get.asserting(_ should contain theSameElementsAs (preLoad ++ postLoad)))
-        _          <- fiber.cancel
+        assertion  <- loadAndRunR(NonEmptyList.one(testTopic1))(
+                        _ => loadState.set(true),
+                        r => topicState.getAndUpdate(_ :+ r).void
+                      ).use { _ =>
+                        for {
+                          _         <- eventually(topicState.get.asserting(_ should contain theSameElementsAs preLoad))
+                          _         <- loadState.get.asserting(_ shouldBe true)
+                          _         <- publishStringMessages(testTopic1, postLoad)
+                          assertion <-
+                            eventually(
+                              topicState.get.asserting(_ should contain theSameElementsAs (preLoad ++ postLoad))
+                            )
+                        } yield assertion
+                      }
       } yield assertion
     }
 

diff --git a/src/test/scala/utils/KafkaHelpers.scala b/src/test/scala/utils/KafkaHelpers.scala
@@ -4,6 +4,8 @@ import java.util.UUID
 
 import base.AsyncIntSpec
 import cats.data.{NonEmptyList, NonEmptySet}
+import cats.effect.implicits.*
+import cats.effect.kernel.Fiber
 import cats.effect.{Async, Resource}
 import cats.syntax.all.*
 import fs2.Stream
@@ -66,6 +68,21 @@ trait KafkaHelpers[F[_]] {
     TopicLoader.load(topics, strategy, consumerSettings).compile.toList.map(_.map(recordToTuple))
   }
 
+  def loadAndRunR(topics: NonEmptyList[String])(
+      onLoad: Resource.ExitCase => F[Unit],
+      onRecord: ((String, String)) => F[Unit]
+  )(implicit
+      consumerSettings: ConsumerSettings[F, String, String],
+      F: Async[F]
+  ): Resource[F, Fiber[F, Throwable, Unit]] = Resource.make {
+    loadAndRunLoader(topics)(onLoad)
+      .map(recordToTuple)
+      .evalTap(onRecord)
+      .compile
+      .drain
+      .start
+  }(_.cancel.void)
+
   def loadAndRunLoader(topics: NonEmptyList[String])(onLoad: Resource.ExitCase => F[Unit])(implicit
       consumerSettings: ConsumerSettings[F, String, String],
       F: Async[F]