Added async seed support to harness-zio-mock

modules/harness-zio-mock/shared/src/main/scala/harness/zio/mock/Mock.scala

@@ -1,5 +1,6 @@
 package harness.zio.mock
 
+import harness.zio.mock.error.MockError
 import java.util.UUID
 import zio.*
 
@@ -102,6 +103,31 @@ abstract class Mock[Z](implicit rTag: Tag[Z]) { myMock =>
       def failure(f: => E): URIO[Proxy & Z, Unit] =
         zioI(_ => ZIO.fail(f))
 
+      def expect(test: PartialFunction[I, Boolean])(implicit ev: Unit <:< A): URIO[Proxy & Z, Unit] =
+        zioI { i => ZIO.die(MockError.FailedExpectedSeed(myCapability, i)).unless(test.lift(i).getOrElse(false)).as(ev(())) }
+
+      object async {
+
+        def zioI(f: I => ZIO[R, E, A]): URIO[Proxy & Z, Unit] =
+          ZIO.serviceWithZIO[Proxy](_.appendAsyncSeed(myCapability, f))
+
+        def zio(f: => ZIO[R, E, A]): URIO[Proxy & Z, Unit] =
+          zioI(_ => f)
+
+        def successI(f: I => A): URIO[Proxy & Z, Unit] =
+          zioI(i => ZIO.succeed(f(i)))
+
+        def success(f: => A): URIO[Proxy & Z, Unit] =
+          zioI(_ => ZIO.succeed(f))
+
+        def failureI(f: I => E): URIO[Proxy & Z, Unit] =
+          zioI(i => ZIO.fail(f(i)))
+
+        def failure(f: => E): URIO[Proxy & Z, Unit] =
+          zioI(_ => ZIO.fail(f))
+
+      }
+
       object prepend {
 
         def zioI(f: I => ZIO[R, E, A]): URIO[Proxy & Z, Unit] =

modules/harness-zio-mock/shared/src/main/scala/harness/zio/mock/Proxy.scala

@@ -7,19 +7,30 @@ import zio.*
 
 final class Proxy private (
     private val implsRef: Ref[Map[ErasedCapability, ErasedEffectImpl]],
-    private val expectationsRef: Ref[List[ErasedExpectation]],
+    private val expectationsRef: Ref[List[Proxy.Seed]],
 ) {
 
   def apply[Z, I, R, E, A](
       capability: Mock[Z]#Capability[I, R, E, A],
       i: I,
-  ): ZIO[R, E, A] =
+  ): ZIO[R, E, A] = {
+    object forCapability {
+      def unapply(map: Map[ErasedCapability, ErasedExpectation]): Option[(EffectImpl[I, R, E, A], Option[Map[ErasedCapability, ErasedExpectation]])] =
+        map.get(capability).map { exp =>
+          val effect: EffectImpl[I, R, E, A] = exp.effect.asInstanceOf
+          val newMap = map.removed(capability)
+          (effect, Option.when(newMap.nonEmpty)(newMap))
+        }
+    }
+
     for {
       impl <- implsRef.get.map(_.get(capability))
       effect <- expectationsRef.modify {
-        case head :: tail if head.capability == capability =>
+        case Proxy.Seed.Sync(head) :: tail if head.capability == capability =>
           val result: EffectImpl[I, R, E, A] = head.effect.asInstanceOf
           (result, tail)
+        case Proxy.Seed.Async(forCapability(result, newExps)) :: tail =>
+          (result, newExps.fold(tail)(exps => Proxy.Seed.Async(exps) :: tail))
         case seededExpectations =>
           impl match {
             case Some(effectImpl) =>
@@ -30,13 +41,14 @@ final class Proxy private (
               )
             case None =>
               (
-                (_: I) => ZIO.die(MockError.UnexpectedCall(capability, seededExpectations.map(_.capability))),
+                (_: I) => ZIO.die(MockError.UnexpectedCall(capability, seededExpectations)),
                 seededExpectations,
               )
           }
       }
       result <- effect(i)
     } yield result
+  }
 
   def apply[Z, I, R, E, A](
       capability: Mock[Z]#Capability[I, R, E, A],
@@ -149,25 +161,41 @@ final class Proxy private (
       implsRef.update(_.updated(capability, effect))
 
   private[mock] def prependSeed[Z, I, R, E, A](capability: Mock[Z]#Capability[I, R, E, A], effect: I => ZIO[R, E, A]): UIO[Unit] =
-    expectationsRef.update(Expectation(capability, effect) :: _)
+    expectationsRef.update(Proxy.Seed.Sync(Expectation(capability, effect)) :: _)
 
   private[mock] def appendSeed[Z, I, R, E, A](capability: Mock[Z]#Capability[I, R, E, A], effect: I => ZIO[R, E, A]): UIO[Unit] =
-    expectationsRef.update(_ :+ Expectation(capability, effect))
+    expectationsRef.update(_ :+ Proxy.Seed.Sync(Expectation(capability, effect)))
+
+  private[mock] def appendAsyncSeed[Z, I, R, E, A](capability: Mock[Z]#Capability[I, R, E, A], effect: I => ZIO[R, E, A]): UIO[Unit] = {
+    val expectation: ErasedExpectation = Expectation(capability, effect)
+    expectationsRef.modify { expectations =>
+      expectations.reverse match {
+        case Proxy.Seed.Async(exps) :: tail =>
+          if (exps.contains(expectation.capability)) (ZIO.die(MockError.DuplicateAsyncSeed(capability)), expectations)
+          else (ZIO.unit, (Proxy.Seed.Async(exps.updated(expectation.capability, expectation)) :: tail).reverse)
+        case rev =>
+          (ZIO.unit, (Proxy.Seed.Async(Map(expectation.capability -> expectation)) :: rev).reverse)
+      }
+    }.flatten
+  }
+
+  private val appendEmptyAsync: UIO[Unit] =
+    expectationsRef.update(_ :+ Proxy.Seed.Async(Map.empty))
 
 }
 object Proxy {
 
   private val makeWithoutFinalizer: UIO[Proxy] =
     for {
       implsRef <- Ref.make(Map.empty[ErasedCapability, ErasedEffectImpl])
-      expectationsRef <- Ref.make(List.empty[ErasedExpectation])
+      expectationsRef <- Ref.make(List.empty[Seed])
     } yield new Proxy(implsRef, expectationsRef)
 
   private val makeWithFinalizer: URIO[Scope, Proxy] =
     makeWithoutFinalizer.withFinalizer {
       _.expectationsRef.get.flatMap {
         _.toNel match {
-          case Some(unsatisfiedExpectations) => ZIO.die(MockError.UnsatisfiedCalls(unsatisfiedExpectations.map(_.capability)))
+          case Some(unsatisfiedExpectations) => ZIO.die(MockError.UnsatisfiedCalls(unsatisfiedExpectations))
           case None                          => ZIO.unit
         }
       }
@@ -176,4 +204,25 @@ object Proxy {
   val layer: ULayer[Proxy] =
     ZLayer.scoped { makeWithFinalizer }
 
+  val seedEmptyAsync: URIO[Proxy, Unit] =
+    ZIO.serviceWithZIO[Proxy](_.appendEmptyAsync)
+
+  sealed trait Seed {
+
+    final def capabilities: Set[ErasedCapability] = this match {
+      case Seed.Sync(expectation)   => Set(expectation.capability)
+      case Seed.Async(expectations) => expectations.keySet
+    }
+
+    final def hasCapability(capability: ErasedCapability): Boolean = this match {
+      case Seed.Sync(expectation)   => expectation.capability == capability
+      case Seed.Async(expectations) => expectations.contains(capability)
+    }
+
+  }
+  object Seed {
+    final case class Sync(expectation: ErasedExpectation) extends Seed
+    final case class Async(expectations: Map[ErasedCapability, ErasedExpectation]) extends Seed
+  }
+
 }

modules/harness-zio-mock/shared/src/main/scala/harness/zio/mock/error/MockError.scala

@@ -1,23 +1,33 @@
 package harness.zio.mock.error
 
 import cats.data.NonEmptyList
+import cats.syntax.option.*
 import harness.zio.mock.*
+import harness.zio.mock.Proxy.Seed
 import harness.zio.mock.Types.*
 
 sealed trait MockError extends Throwable {
   import MockError.{attnStar, hintHeader}
 
   // TODO (KR) : colorize?
   override final def getMessage: String = {
+    def showSeeds(seeds: List[Proxy.Seed], attn: Option[ErasedCapability]): String =
+      seeds.map {
+        case Seed.Sync(expectation) =>
+          s"\n    - ${if (attn.contains(expectation.capability)) attnStar else ""}${expectation.capability.name}"
+        case Seed.Async(expectations) =>
+          s"\n    - [Async]:${expectations.values.map(e => s"\n      - ${if (attn.contains(e.capability)) attnStar else ""}${e.capability.name}").mkString}"
+      }.mkString
+
     val baseMessage: String =
       this match {
         case MockError.UnexpectedCall(givenCapability, expectations) =>
           val seededCallsStr = expectations match {
             case Nil => "\n    (expected no more calls)"
-            case _   => expectations.map(e => s"\n    - ${if (e == givenCapability) attnStar else ""}${e.name}").mkString
+            case _   => showSeeds(expectations, givenCapability.some)
           }
 
-          expectations.indexOf(givenCapability) match {
+          expectations.indexWhere(_.hasCapability(givenCapability)) match {
             case -1 =>
               s"""Unexpected call to capability ${givenCapability.name}
                  |  currently seeded calls:$seededCallsStr
@@ -34,7 +44,7 @@ sealed trait MockError extends Throwable {
                  |  currently seeded calls:$seededCallsStr
                  |$hintHeader
                  |  It looks like you seeded a call to ${givenCapability.name} later on.
-                 |  Calls that are expected before that call: ${unexpected.map(e => s"\n    - ${e.name}").mkString}
+                 |  Calls that are expected before that call:${showSeeds(unexpected, None)}
                  |  Likely causes & solutions:
                  |    1) You do in deed expect these other calls to happen first.
                  |       solution: fix your code such that these other calls happen first.
@@ -44,7 +54,7 @@ sealed trait MockError extends Throwable {
                  |       solution: your code that is seeding calls is incorrect, remove the code that seeds these other expectations.""".stripMargin
           }
         case MockError.UnsatisfiedCalls(expectations) =>
-          s"""Unsatisfied seeded calls:${expectations.toList.map(e => s"\n  - ${e.name}").mkString}
+          s"""Unsatisfied seeded calls:${showSeeds(expectations.toList, None)}
              |$hintHeader
              |  Tests can not exit with seeded calls that did not end up being called.
              |  Likely causes & solutions:
@@ -58,6 +68,19 @@ sealed trait MockError extends Throwable {
              |$hintHeader
              |  You can not implement the same capability multiple times.
              |  You essentially did $mockImplStr ++ $mockImplStr""".stripMargin
+        case MockError.FailedExpectedSeed(capability, input) =>
+          val showInput = (input.asInstanceOf[Matchable] match {
+            case tup: Tuple => tup.toList
+            case _          => input :: Nil
+          }).map(e => s"\n    - $e").mkString
+          s"""Capability called with invalid input: ${capability.name}
+             |  actual input:$showInput
+             |$hintHeader
+             |  Either you have an error in your code, or your seeded expectation is asserting the wrong thing.""".stripMargin
+        case MockError.DuplicateAsyncSeed(capability) =>
+          s"""Duplicate async capability seeded: ${capability.name}
+             |$hintHeader
+             |  It's not supported to seed multiple async expectations for the same capability.""".stripMargin
       }
 
     s"\n$baseMessage".replaceAll("\n", "\n  ")
@@ -73,15 +96,24 @@ object MockError {
 
   final case class UnexpectedCall(
       givenCapability: ErasedCapability,
-      expectations: List[ErasedCapability],
+      expectations: List[Proxy.Seed],
   ) extends MockError
 
   final case class UnsatisfiedCalls(
-      expectations: NonEmptyList[ErasedCapability],
+      expectations: NonEmptyList[Proxy.Seed],
   ) extends MockError
 
   final case class CapabilityIsAlreadyImplemented(
       capability: ErasedCapability,
   ) extends MockError
 
+  final case class FailedExpectedSeed(
+      capability: ErasedCapability,
+      input: Any,
+  ) extends MockError
+
+  final case class DuplicateAsyncSeed(
+      capability: ErasedCapability,
+  ) extends MockError
+
 }

modules/harness-zio-mock/shared/src/test/scala/harness/zio/mock/MockSpec.scala

@@ -1,6 +1,5 @@
 package harness.zio.mock
 
-import cats.data.NonEmptyList
 import harness.zio.*
 import harness.zio.mock.error.MockError
 import harness.zio.test.*
@@ -13,25 +12,31 @@ object MockSpec extends DefaultHarnessSpec {
   private trait ExService {
     def funct1(i: Int): IO[String, Int]
     def funct2(i1: Int, i2: Double): IO[String, Double]
+    def funct3(s: String): IO[String, Int]
   }
   private object ExService {
     def funct1(i: Int): ZIO[ExService, String, Int] =
       ZIO.serviceWithZIO[ExService](_.funct1(i))
     def funct2(i1: Int, i2: Double): ZIO[ExService, String, Double] =
       ZIO.serviceWithZIO[ExService](_.funct2(i1, i2))
+    def funct3(s: String): ZIO[ExService, String, Int] =
+      ZIO.serviceWithZIO[ExService](_.funct3(s))
   }
 
   private object ExServiceMock extends Mock[ExService] {
 
     object Funct1 extends Effect[Int, String, Int]
     object Funct2 extends Effect[(Int, Double), String, Double]
+    object Funct3 extends Effect[String, String, Int]
 
     override protected def buildInternal(proxy: Proxy): ExService =
       new ExService {
         override def funct1(i: Int): IO[String, Int] =
           proxy(Funct1, i)
         override def funct2(i1: Int, i2: Double): IO[String, Double] =
           proxy(Funct2, i1, i2)
+        override def funct3(s: String): IO[String, Int] =
+          proxy(Funct3, s)
       }
 
   }
@@ -135,7 +140,10 @@ object MockSpec extends DefaultHarnessSpec {
           _ <- ExServiceMock.Funct1.seed.success(1)
           _ <- ExServiceMock.Funct2.seed.success(1)
         } yield assertCompletes
-      } @@ TestAspect.failing(_ == TestFailure.die(MockError.UnsatisfiedCalls(NonEmptyList.of(ExServiceMock.Funct1, ExServiceMock.Funct2)))),
+      } @@ TestAspect.failing {
+        case TestFailure.Runtime(Cause.Die(_: MockError.UnsatisfiedCalls, _), _) => true
+        case _                                                                   => false
+      },
       makeTest("fails if there are no seeds") {
         ExServiceMock.empty
       } {
@@ -150,7 +158,7 @@ object MockSpec extends DefaultHarnessSpec {
           _ <- ExServiceMock.Funct2.seed.success(0)
           res <- ExService.funct1(0).exit
           _ <- ExService.funct2(0, 0) // empty seed
-        } yield assert(res)(dies(equalTo(MockError.UnexpectedCall(ExServiceMock.Funct1, List(ExServiceMock.Funct2)))))
+        } yield assert(res)(dies(isSubtype[MockError.UnexpectedCall](anything)))
       },
     )
 
@@ -172,6 +180,46 @@ object MockSpec extends DefaultHarnessSpec {
       },
     )
 
+  private val positiveAsyncSpec: TestSpec =
+    suite("positive")(
+      makeSeedTest("allows out of order calls - 1") {
+        for {
+          _ <- ExServiceMock.Funct1.seed.success(1)
+          _ <- ExServiceMock.Funct2.seed.async.success(2)
+          _ <- ExServiceMock.Funct3.seed.async.success(3)
+          _ <- ExServiceMock.Funct1.seed.success(4)
+
+          res1 <- ExService.funct1(0)
+          res2 <- ExService.funct2(0, 0)
+          res3 <- ExService.funct3("")
+          res4 <- ExService.funct1(0)
+        } yield assertTrue(
+          res1 == 1,
+          res2 == 2,
+          res3 == 3,
+          res4 == 4,
+        )
+      },
+      makeSeedTest("allows out of order calls - 2") {
+        for {
+          _ <- ExServiceMock.Funct1.seed.success(1)
+          _ <- ExServiceMock.Funct2.seed.async.success(2)
+          _ <- ExServiceMock.Funct3.seed.async.success(3)
+          _ <- ExServiceMock.Funct1.seed.success(4)
+
+          res1 <- ExService.funct1(0)
+          res3 <- ExService.funct3("")
+          res2 <- ExService.funct2(0, 0)
+          res4 <- ExService.funct1(0)
+        } yield assertTrue(
+          res1 == 1,
+          res2 == 2,
+          res3 == 3,
+          res4 == 4,
+        )
+      },
+    ).provideSomeLayer[HarnessEnv](Proxy.layer >+> ExServiceMock.empty.toLayer)
+
   override def spec: TestSpec =
     suite("ExServiceSpec")(
       suite("impl")(
@@ -184,6 +232,9 @@ object MockSpec extends DefaultHarnessSpec {
       suite("impl + seed")(
         positiveImplAndSeedSpec,
       ),
+      suite("async seed")(
+        positiveAsyncSpec,
+      ),
     )
 
 }

project/build.properties

@@ -1 +1 @@
-sbt.version=1.6.2
+sbt.version=1.9.9