Tweaked Mocked.toLayer to not make its own proxy

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

@@ -1,71 +1,51 @@
 package harness.zio.mock
 
-import cats.syntax.list.*
 import harness.zio.mock.Types.*
-import harness.zio.mock.error.MockError
 import zio.*
 
-sealed trait Mocked[Z] { self =>
-
-  final def ++[Z2](that: Mocked[Z2]): Mocked[Z & Z2] =
-    (self.asInstanceOf[Mocked[Z & Z2]], that.asInstanceOf[Mocked[Z & Z2]]) match {
-      case (self: Mocked.Building[Z & Z2], that: Mocked.Building[Z & Z2]) =>
-        (self.impls.keySet & that.impls.keySet).toList.toNel match {
-          case None =>
-            new Mocked.Building[Z & Z2](
-              impls = self.impls ++ that.impls,
-              mocks = self.mocks | that.mocks,
-            )
-          case Some(conflictingCapabilities) =>
-            Mocked.FailedDuringBuild(MockError.OverlappingCapabilityImplementations(conflictingCapabilities))
-        }
-      case (self: Mocked.FailedDuringBuild[Z & Z2], _)                          => self
-      case (_: Mocked.Building[Z & Z2], that: Mocked.FailedDuringBuild[Z & Z2]) => that
-    }
-
-  final def toLayer(implicit envTag: EnvironmentTag[Z]): ULayer[Proxy & Z] =
-    self match {
-      case self: Mocked.Building[Z] =>
-        Proxy.make(self.impls.asInstanceOf[Map[ErasedCapability, ErasedEffectImpl]]) >+>
-          ZLayer.fromZIOEnvironment {
-            ZIO.serviceWith[Proxy] { proxy =>
-              val zEnvs: List[ZEnvironment[?]] =
-                self.mocks.toList.map { _.build(proxy) }
-              val zEnv: ZEnvironment[Z] =
-                zEnvs.foldLeft(ZEnvironment.empty) { _ ++ _ }.asInstanceOf[ZEnvironment[Z]]
-
-              zEnv
-            }
-          }
-      case Mocked.FailedDuringBuild(error) =>
-        ZLayer.die(error)
+final class Mocked[Z](
+    private[Mocked] val impls: List[Expectation[? >: Z, ?, ?, ?, ?]],
+    private[Mocked] val mocks: Set[Mock[? >: Z]],
+) { self =>
+
+  final def ++[Z2](that: Mocked[Z2]): Mocked[Z & Z2] = {
+    val _self: Mocked[Z & Z2] = self.asInstanceOf
+    val _that: Mocked[Z & Z2] = that.asInstanceOf
+    new Mocked[Z & Z2](
+      impls = _self.impls ::: _that.impls,
+      mocks = _self.mocks | _that.mocks,
+    )
+  }
+
+  final def toLayer: URLayer[Proxy, Z] =
+    ZLayer.fromZIOEnvironment {
+      for {
+        // get proxy
+        proxy <- ZIO.service[Proxy]
+
+        // add impls to proxy
+        _ <- ZIO.foreachDiscard(impls)(i => proxy.putImpl(i.capability, i.effect))
+
+        // create env
+        zEnvs: List[ZEnvironment[?]] =
+          self.mocks.toList.map { _.build(proxy) }
+        zEnv: ZEnvironment[Z] =
+          zEnvs.foldLeft(ZEnvironment.empty) { _ ++ _ }.asInstanceOf[ZEnvironment[Z]]
+      } yield zEnv
     }
 
 }
 object Mocked {
 
-  // =====| Builders |=====
-
-  private final class Building[Z](
-      private[Mocked] val impls: Map[ErasedCapabilityZ[? >: Z], ErasedEffectImpl],
-      private[Mocked] val mocks: Set[Mock[? >: Z]],
-  ) extends Mocked[Z]
-
-  private final case class FailedDuringBuild[Z](
-      error: Throwable,
-  ) extends Mocked[Z]
-
-  // =====|  |=====
-
   private[mock] def single[Z, I, R, E, A](capability: Mock[Z]#Capability[I, R, E, A], effect: I => ZIO[R, E, A]): Mocked[Z] =
-    new Mocked.Building[Z](
-      Map(capability -> effect),
+    new Mocked[Z](
+      Expectation(capability, effect) :: Nil,
       Set(capability.getMock),
     )
 
   private[mock] def empty[Z](mock: Mock[Z]): Mocked[Z] =
-    new Mocked.Building[Z](
-      Map.empty,
+    new Mocked[Z](
+      Nil,
       Set(mock),
     )
 

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

@@ -6,15 +6,15 @@ import harness.zio.mock.error.MockError
 import zio.*
 
 final class Proxy private (
-    private val impls: Map[ErasedCapability, ErasedEffectImpl],
+    private val implsRef: Ref[Map[ErasedCapability, ErasedEffectImpl]],
     private val expectationsRef: Ref[List[ErasedExpectation]],
 ) {
 
   def apply[Z, I, R, E, A](
       capability: Mock[Z]#Capability[I, R, E, A],
       i: I,
   ): ZIO[R, E, A] =
-    impls.get(capability) match {
+    implsRef.get.map(_.get(capability)).flatMap {
       case Some(effectImpl) =>
         val typed: EffectImpl[I, R, E, A] = effectImpl.asInstanceOf
         typed(i)
@@ -139,19 +139,33 @@ 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) *>
+      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] =
-    ZIO.die(MockError.CanNotSeedImplementedCapability(capability)).when(impls.contains(capability)) *>
+    ensureCapabilityIsNotImplemented(capability) *>
       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] =
-    ZIO.die(MockError.CanNotSeedImplementedCapability(capability)).when(impls.contains(capability)) *>
+    ensureCapabilityIsNotImplemented(capability) *>
       expectationsRef.update(_ :+ Expectation(capability, effect))
 
 }
 object Proxy {
 
-  private def makeScoped(impls: Map[ErasedCapability, ErasedEffectImpl]): URIO[Scope, Proxy] =
-    Ref.make(List.empty[ErasedExpectation]).map { new Proxy(impls, _) }.withFinalizer {
+  private val makeWithoutFinalizer: UIO[Proxy] =
+    for {
+      implsRef <- Ref.make(Map.empty[ErasedCapability, ErasedEffectImpl])
+      expectationsRef <- Ref.make(List.empty[ErasedExpectation])
+    } 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)))
@@ -160,7 +174,7 @@ object Proxy {
       }
     }
 
-  private[mock] def make(impls: Map[ErasedCapability, ErasedEffectImpl]): ULayer[Proxy] =
-    ZLayer.scoped { makeScoped(impls) }
+  val layer: ULayer[Proxy] =
+    ZLayer.scoped { makeWithFinalizer }
 
 }

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

@@ -16,10 +16,8 @@ sealed trait MockError extends Throwable {
       s"\n  Unexpected call to capability ${givenCapability.name}$expStr"
     case MockError.UnsatisfiedCalls(expectations) =>
       s"\n  Unsatisfied seeded calls (test can not exit with remaining expectations):${expectations.toList.map(e => s"\n    - ${e.name}").mkString}"
-    case MockError.OverlappingCapabilityImplementations(capabilities) =>
-      s"\n  Conflicting capabilities provided for:${capabilities.toList.map(c => s"\n    - ${c.name}").mkString}"
-    case MockError.CanNotSeedImplementedCapability(capability) =>
-      s"\n  Can not seed capability which is implemented: ${capability.name}"
+    case MockError.CapabilityIsAlreadyImplemented(capability) =>
+      s"\n  Capability is already implemented: ${capability.name}"
   }
 
   override final def toString: String = getMessage
@@ -36,11 +34,7 @@ object MockError {
       expectations: NonEmptyList[ErasedCapability],
   ) extends MockError
 
-  final case class OverlappingCapabilityImplementations(
-      capabilities: NonEmptyList[ErasedCapability],
-  ) extends MockError
-
-  final case class CanNotSeedImplementedCapability(
+  final case class CapabilityIsAlreadyImplemented(
       capability: ErasedCapability,
   ) extends MockError
 

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

@@ -1,8 +1,8 @@
 package harness.zio.mock
 
 import harness.zio.*
-import harness.zio.test.*
 import harness.zio.mock.*
+import harness.zio.test.*
 import java.util.UUID
 import zio.*
 import zio.test.*
@@ -97,6 +97,7 @@ object ExampleSpec extends DefaultHarnessSpec {
 
   override def spec: TestSpec =
     specImpl.provideSome[HarnessEnv](
+      Proxy.layer,
       (
         Service1Mock.GetPerson.implement.success(Person(UUID.randomUUID, "first", "last", 18)) ++
           Service2Mock.empty

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

@@ -2,8 +2,8 @@ package harness.zio.mock
 
 import cats.data.NonEmptyList
 import harness.zio.*
-import harness.zio.test.*
 import harness.zio.mock.error.MockError
+import harness.zio.test.*
 import zio.*
 import zio.test.*
 import zio.test.Assertion.*
@@ -37,7 +37,7 @@ object MockSpec extends DefaultHarnessSpec {
   }
 
   private def makeTest(name: String)(mocked: Mocked[ExService])(testF: => ZIO[HarnessEnv & ExService & Proxy, Any, TestResult]): TestSpec =
-    test(name) { testF.provideSomeLayer[HarnessEnv](mocked.toLayer) }
+    test(name) { testF.provideSomeLayer[HarnessEnv](Proxy.layer >+> mocked.toLayer) }
 
   private def makeSeedTest(name: String)(testF: => ZIO[HarnessEnv & ExService & Proxy, Any, TestResult]): Spec[HarnessEnv & ExService & Proxy, Any] =
     test(name) { testF }
@@ -124,7 +124,7 @@ object MockSpec extends DefaultHarnessSpec {
           res4 == 125,
         )
       },
-    ).provideSomeLayer[HarnessEnv](ExServiceMock.empty.toLayer)
+    ).provideSomeLayer[HarnessEnv](Proxy.layer >+> ExServiceMock.empty.toLayer)
 
   private val negativeSeedSpec: TestSpec =
     suite("negative")(