UnsatisfiedExpectationsException when mocks used as layers

[michalkoza]

I use Scala 3.3.3, ZIO 2.1.9, zio-mock: 1.0.0-RC12

Scastie: https://scastie.scala-lang.org/ICPrp4mkQJCBieDRi3gsiA

I have this trait:

trait MyTrait {
  def myMethod(myArg: String): ZIO[Any, Nothing, Int]
}

a mock for it:

import zio.*
import zio.mock.*

object MyTraitMock extends Mock[MyTrait] {
  object MyMethod extends Effect[String, Nothing, Int]

  val compose: URLayer[Proxy, MyTrait] =
    ZLayer {
      ZIO.serviceWith[Proxy] { proxy => (myArg: String) =>
        for {
          res <- proxy(MyMethod, myArg)
          _ <- ZIO.succeed(println(s"CALL, myArg: $myArg, res: $res"))
        } yield res
      }
    }
}

and a test suite with a buch of almost identical tests.
They should all pass, but only the first one and the last one pass. Others fail as if mock were never called, but it must have been.

import zio.mock.Expectation
import zio.test.Assertion.equalTo
import zio.test.{Spec, TestEnvironment, ZIOSpecDefault, assertTrue}
import zio.{Scope, ZIO}

object MyTraitSpec extends ZIOSpecDefault {
  override def spec: Spec[TestEnvironment & Scope, Any] =
    suite("suite")(
      test("test with atMost") {
        val layer = MyTraitMock.MyMethod(assertion = equalTo("arg1"), result = Expectation.value(10)).atMost(1).toLayer
        for {
          mockInstance <- ZIO.service[MyTrait].provideLayer(layer)
          result1 <- mockInstance.myMethod("arg1")
        } yield assertTrue(
          result1 == 10,
        )
      },

      test("test without quantifier") {
        val layer = MyTraitMock.MyMethod(assertion = equalTo("arg1"), result = Expectation.value(10)).toLayer
        for {
          mockInstance <- ZIO.service[MyTrait].provideLayer(layer)
          result1 <- mockInstance.myMethod("arg1")
        } yield assertTrue(
          result1 == 10,
        )
      },

      test("test with atLeast") {
        val layer = MyTraitMock.MyMethod(assertion = equalTo("arg1"), result = Expectation.value(10)).atLeast(1).toLayer
        for {
          mockInstance <- ZIO.service[MyTrait].provideLayer(layer)
          result1 <- mockInstance.myMethod("arg1")
        } yield assertTrue(
          result1 == 10,
        )
      },

      test("test with exactly") {
        val layer = MyTraitMock.MyMethod(assertion = equalTo("arg1"), result = Expectation.value(10)).exactly(2).toLayer
        for {
          mockInstance <- ZIO.service[MyTrait].provideLayer(layer)
          result1 <- mockInstance.myMethod("arg1")
          result2 <- mockInstance.myMethod("arg1")
        } yield assertTrue(
          result1 == 10,
          result2 == 10,
        )
      },

      test("test with exactly") {
        for {
          env <- MyTraitMock.MyMethod(assertion = equalTo("arg1"), result = Expectation.value(10)).exactly(2).toLayer.build
          mockInstance = env.get[MyTrait]
          result1 <- mockInstance.myMethod("arg1")
          result2 <- mockInstance.myMethod("arg1")
        } yield assertTrue(
          result1 == 10,
          result2 == 10,
        )
      }
    )
}

I think there is some issue with expectation count quantifiers when mocks are used as layers.