diff --git a/modules/core/shared/src/test/scala/KleisliTest.scala b/modules/core/shared/src/test/scala/SpanPropagationTest.scala similarity index 62% rename from modules/core/shared/src/test/scala/KleisliTest.scala rename to modules/core/shared/src/test/scala/SpanPropagationTest.scala index b96a3fa8..74070cb9 100644 --- a/modules/core/shared/src/test/scala/KleisliTest.scala +++ b/modules/core/shared/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)) } }