change InMemorySuite.LocalSuite to more closely test the actual scena…

…rio being contemplated
typelevel · Jan 31, 2023 · b32fbed · b32fbed
1 parent abab2f1
commit b32fbed
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 54 deletions.
diff --git a/modules/core-tests/shared/src/test/scala/InMemorySuite.scala b/modules/core-tests/shared/src/test/scala/InMemorySuite.scala
@@ -4,10 +4,11 @@
 
 package natchez
 
-import cats.Applicative
+import cats.arrow.FunctionK
+import cats.{Applicative, ~>}
 import cats.syntax.all._
 import cats.data.Kleisli
-import cats.effect.{IO, IOLocal, MonadCancelThrow}
+import cats.effect.{Concurrent, IO, IOLocal, MonadCancelThrow}
 import cats.mtl.Local
 import munit.CatsEffectSuite
 import natchez.InMemory.Lineage.defaultRootName
@@ -67,74 +68,67 @@ object InMemorySuite {
   }
 
   trait LocalSuite extends InMemorySuite {
+    type LocalProgram[F[_]] = (EntryPoint[F], Local[F, Span[F]]) => F[Unit]
+
     trait LocalTest {
-      def program[F[_]: MonadCancelThrow](implicit L: Local[F, Span[F]]): F[Unit]
+      def program[F[_]: MonadCancelThrow](entryPoint: EntryPoint[F])(implicit
+          L: Local[F, Span[F]]
+      ): F[Unit]
       def expectedHistory: List[(Lineage, NatchezCommand)]
     }
 
     def localTest(name: String, tt: LocalTest): Unit = {
       test(s"$name - Kleisli")(
-        testLocalKleisli(tt.program[Kleisli[IO, Span[IO], *]](implicitly, _), tt.expectedHistory)
+        testLocalKleisli(tt.program[Kleisli[IO, Span[IO], *]](_)(implicitly, _), tt.expectedHistory)
+      )
+      test(s"$name - IOLocal")(
+        testLocalIoLocal(tt.program[IO](_)(implicitly, _), tt.expectedHistory)
       )
-      test(s"$name - IOLocal")(testLocalIoLocal(tt.program[IO](implicitly, _), tt.expectedHistory))
     }
 
-    def testLocalKleisli(
-        localProgram: Local[Kleisli[IO, Span[IO], *], Span[Kleisli[IO, Span[IO], *]]] => Kleisli[
-          IO,
-          Span[IO],
-          Unit
-        ],
+    private def testLocalKleisli(
+        localProgram: LocalProgram[Kleisli[IO, Span[IO], *]],
         expectedHistory: List[(Lineage, NatchezCommand)]
-    ): IO[Unit] = testTraceViaLocal[Kleisli[IO, Span[IO], *]](
-      localProgram,
-      root =>
-        IO.pure(
-          implicitly[Local[Kleisli[IO, Span[IO], *], Span[Kleisli[IO, Span[IO], *]]]] -> (k =>
-            k.run(root)
-          )
-        ),
+    ): IO[Unit] = testProgramGivenEntryPoint[Kleisli[IO, Span[IO], *]](
+      localProgram(_, implicitly),
+      Kleisli.applyK(Span.noop[IO]),
       expectedHistory
     )
 
-    def testLocalIoLocal(
-        localProgram: Local[IO, Span[IO]] => IO[Unit],
+    private def testLocalIoLocal(
+        localProgram: LocalProgram[IO],
         expectedHistory: List[(Lineage, NatchezCommand)]
     ): IO[Unit] =
-      testTraceViaLocal[IO](
-        localProgram,
-        IOLocal(_)
-          .map { ioLocal =>
-            new Local[IO, Span[IO]] {
-              override def local[A](fa: IO[A])(f: Span[IO] => Span[IO]): IO[A] =
-                ioLocal.get.flatMap { initial =>
-                  ioLocal.set(f(initial)) >> fa.guarantee(ioLocal.set(initial))
-                }
-
-              override def applicative: Applicative[IO] = implicitly
-
-              override def ask[E2 >: Span[IO]]: IO[E2] = ioLocal.get
+      testProgramGivenEntryPoint[IO](
+        ep =>
+          IOLocal(Span.noop[IO])
+            .map { ioLocal =>
+              new Local[IO, Span[IO]] {
+                override def local[A](fa: IO[A])(f: Span[IO] => Span[IO]): IO[A] =
+                  ioLocal.get.flatMap { initial =>
+                    ioLocal.set(f(initial)) >> fa.guarantee(ioLocal.set(initial))
+                  }
+
+                override def applicative: Applicative[IO] = implicitly
+
+                override def ask[E2 >: Span[IO]]: IO[E2] = ioLocal.get
+              }
             }
-          }
-          .tupleRight(identity[IO[Unit]]),
+            .flatMap(localProgram(ep, _)),
+        FunctionK.id,
         expectedHistory
       )
 
-    def testTraceViaLocal[F[_]](
-        localProgram: Local[F, Span[F]] => F[Unit],
-        makeTraceAndResolver: Span[IO] => IO[(Local[F, Span[F]], F[Unit] => IO[Unit])],
+    private def testProgramGivenEntryPoint[F[_]: Concurrent](
+        localProgram: EntryPoint[F] => F[Unit],
+        fk: F ~> IO,
         expectedHistory: List[(Lineage, NatchezCommand)]
     ): IO[Unit] =
-      InMemory.EntryPoint.create[IO].flatMap { ep =>
-        val traced = ep.root(defaultRootName).use { r =>
-          makeTraceAndResolver(r).flatMap { case (localInstance, resolve) =>
-            resolve(localProgram(localInstance))
-          }
-        }
-        traced *> ep.ref.get.map { history =>
+      fk(InMemory.EntryPoint.create[F].flatMap { ep =>
+        localProgram(ep) *> ep.ref.get.map { history =>
           assertEquals(history.toList, expectedHistory)
         }
-      }
+      })
   }
 
 }
diff --git a/modules/core-tests/shared/src/test/scala/LocalTraceSpec.scala b/modules/core-tests/shared/src/test/scala/LocalTraceSpec.scala
@@ -7,23 +7,26 @@ package mtl
 
 import cats.effect.{MonadCancelThrow, Trace => _}
 import cats.mtl._
-import natchez.InMemory.Lineage.defaultRootName
 
 class LocalTraceSpec extends InMemorySuite.LocalSuite {
 
   private def useTrace[F[_]: Trace]: F[Unit] = Trace[F].log("hello world")
 
   localTest(
-    "should compile with",
+    "Trace[F] should use root span generated by provided EntryPoint[F] via Local[F, Span[F]]",
     new LocalTest {
-      override def program[F[_]: MonadCancelThrow](implicit L: Local[F, Span[F]]): F[Unit] =
-        useTrace[F]
+      override def program[F[_]: MonadCancelThrow](entryPoint: EntryPoint[F])(implicit
+          L: Local[F, Span[F]]
+      ): F[Unit] =
+        entryPoint
+          .root("my root")
+          .use(Local[F, Span[F]].scope(useTrace[F]))
 
       override def expectedHistory: List[(Lineage, NatchezCommand)] = List(
         Lineage.Root -> NatchezCommand
-          .CreateRootSpan(defaultRootName, Kernel(Map()), Span.Options.Defaults),
-        Lineage.Root(defaultRootName) -> NatchezCommand.LogEvent("hello world"),
-        Lineage.Root -> NatchezCommand.ReleaseRootSpan(defaultRootName)
+          .CreateRootSpan("my root", Kernel(Map()), Span.Options.Defaults),
+        Lineage.Root("my root") -> NatchezCommand.LogEvent("hello world"),
+        Lineage.Root -> NatchezCommand.ReleaseRootSpan("my root")
       )
     }
   )