Merge pull request #664 from tpolecat/topic/propagation

Refactor test to include IO instance
typelevel · Nov 22, 2022 · eda8857 · eda8857
2 parents db44cc7 + 61a11cb
commit eda8857
Showing 1 changed file with 18 additions and 5 deletions.
diff --git a/...e/shared/src/test/scala/KleisliTest.scala → .../src/test/scala/SpanPropagationTest.scala b/...e/shared/src/test/scala/KleisliTest.scala → .../src/test/scala/SpanPropagationTest.scala
@@ -10,13 +10,17 @@ import munit.CatsEffectSuite
 
 import InMemory.{Lineage, NatchezCommand}
 
-class KleisliTest extends CatsEffectSuite {
-  test("span propagation") {
-    def prg[F[_]: Trace] =
-      Trace[F].span("parent")(Trace[F].span("child")(Trace[F].put("answer" -> 42)))
+class SpanPropagationTest extends CatsEffectSuite {
+  def prg[F[_]: Trace] =
+    Trace[F].span("parent")(Trace[F].span("child")(Trace[F].put("answer" -> 42)))
 
+  def testPropagation[F[_]](f: Span[IO] => IO[(Trace[F], F[Unit] => IO[Unit])]) =
     InMemory.EntryPoint.create.flatMap { ep =>
-      val traced = ep.root("root").use(prg[Kleisli[IO, Span[IO], *]].run)
+      val traced = ep.root("root").use { r =>
+        f(r).flatMap { case (traceInstance, resolve) =>
+          resolve(prg(traceInstance))
+        }
+      }
       traced *> ep.ref.get.map { history =>
         assertEquals(
           history.toList,
@@ -32,5 +36,14 @@ class KleisliTest extends CatsEffectSuite {
         )
       }
     }
+
+  test("kleisli") {
+    testPropagation[Kleisli[IO, Span[IO], *]](root =>
+      IO.pure(Trace[Kleisli[IO, Span[IO], *]] -> (k => k.run(root)))
+    )
+  }
+
+  test("io") {
+    testPropagation[IO](root => Trace.ioTrace(root).map(_ -> identity))
   }
 }