Allow mock to mix implement/seed, and prefer seeded expectation if th…

…at is the next expected call.
Kalin-Rudnicki · Mar 21, 2024 · 60d91ff · 60d91ff
1 parent b22e9a5
commit 60d91ff
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 23 deletions.
diff --git a/modules/harness-zio-mock/shared/src/main/scala/harness/zio/mock/Proxy.scala b/modules/harness-zio-mock/shared/src/main/scala/harness/zio/mock/Proxy.scala
@@ -14,24 +14,29 @@ final class Proxy private (
       capability: Mock[Z]#Capability[I, R, E, A],
       i: I,
   ): ZIO[R, E, A] =
-    implsRef.get.map(_.get(capability)).flatMap {
-      case Some(effectImpl) =>
-        val typed: EffectImpl[I, R, E, A] = effectImpl.asInstanceOf
-        typed(i)
-      case None =>
-        expectationsRef
-          .modify {
-            case head :: tail if head.capability == capability =>
-              val result: EffectImpl[I, R, E, A] = head.effect.asInstanceOf
-              (result, tail)
-            case seededExpectations =>
+    for {
+      impl <- implsRef.get.map(_.get(capability))
+      effect <- expectationsRef.modify {
+        case head :: tail if head.capability == capability =>
+          val result: EffectImpl[I, R, E, A] = head.effect.asInstanceOf
+          (result, tail)
+        case seededExpectations =>
+          impl match {
+            case Some(effectImpl) =>
+              val typed: EffectImpl[I, R, E, A] = effectImpl.asInstanceOf
+              (
+                typed,
+                seededExpectations,
+              )
+            case None =>
               (
-                _ => ZIO.die(MockError.UnexpectedCall(capability, seededExpectations.map(_.capability))),
+                (_: I) => ZIO.die(MockError.UnexpectedCall(capability, seededExpectations.map(_.capability))),
                 seededExpectations,
               )
           }
-          .flatMap(_(i))
-    }
+      }
+      result <- effect(i)
+    } yield result
 
   def apply[Z, I, R, E, A](
       capability: Mock[Z]#Capability[I, R, E, A],
@@ -139,21 +144,15 @@ final class Proxy private (
 
   // =====|  |=====
 
-  private def ensureCapabilityIsNotImplemented(capability: ErasedCapability): UIO[Unit] =
-    ZIO.die(MockError.CapabilityIsAlreadyImplemented(capability)).whenZIO(implsRef.get.map(_.contains(capability))).unit
-
-  // TODO (KR) : check that capability is not in seeded expectations
   private[mock] def putImpl[Z, I, R, E, A](capability: Mock[Z]#Capability[I, R, E, A], effect: I => ZIO[R, E, A]): UIO[Unit] =
-    ensureCapabilityIsNotImplemented(capability) *>
+    ZIO.die(MockError.CapabilityIsAlreadyImplemented(capability)).whenZIO(implsRef.get.map(_.contains(capability))) *>
       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] =
-    ensureCapabilityIsNotImplemented(capability) *>
-      expectationsRef.update(Expectation(capability, effect) :: _)
+    expectationsRef.update(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] =
-    ensureCapabilityIsNotImplemented(capability) *>
-      expectationsRef.update(_ :+ Expectation(capability, effect))
+    expectationsRef.update(_ :+ Expectation(capability, effect))
 
 }
 object Proxy {

diff --git a/modules/harness-zio-mock/shared/src/test/scala/harness/zio/mock/MockSpec.scala b/modules/harness-zio-mock/shared/src/test/scala/harness/zio/mock/MockSpec.scala
@@ -154,6 +154,24 @@ object MockSpec extends DefaultHarnessSpec {
       },
     )
 
+  private val positiveImplAndSeedSpec: TestSpec =
+    suite("positive")(
+      makeTest("Prefers seeded expectation to impl if that is the next expected call") {
+        ExServiceMock.Funct1.implement.successI { _ * 2 }
+      } {
+        for {
+          res1 <- ExService.funct1(1)
+          _ <- ExServiceMock.Funct1.seed.success(10)
+          res2 <- ExService.funct1(2)
+          res3 <- ExService.funct1(3)
+        } yield assertTrue(
+          res1 == 2,
+          res2 == 10,
+          res3 == 6,
+        )
+      },
+    )
+
   override def spec: TestSpec =
     suite("ExServiceSpec")(
       suite("impl")(
@@ -163,6 +181,9 @@ object MockSpec extends DefaultHarnessSpec {
         positiveSeedSpec,
         negativeSeedSpec,
       ),
+      suite("impl + seed")(
+        positiveImplAndSeedSpec,
+      ),
     )
 
 }