scalaunit-testingmockitoziozio-test

How to correctly verify scheduled invocations in ZIO Test


I am new to ZIO and ZIO Test and I'd like to test a scheduling service I wrote under ZIO v1.0.0RC17:

The service(s):

import zio.{RIO, Schedule}
import zio.clock.Clock
import zio.duration._

trait ModuleA {
  def moduleA: ModuleA.Service
}

object ModuleA {
  trait Service {
    def schedule(arg: Int): RIO[Clock, Unit]
  }
}

trait ModuleALive extends ModuleA {

  def moduleB: ModuleB.Service

  override def moduleA: ModuleA.Service = new ModuleA.Service {
    override def schedule(arg: Int): RIO[Clock, Unit] = {
      moduleB.run(arg).repeat(Schedule.spaced(1 day)).map(_ => ())
    }
  }
}

trait ModuleB {
  def moduleB: ModuleB.Service
}

object ModuleB {
  trait Service {
    def run(arg: Int): RIO[Clock, Unit]
  }
}

The service of ModuleA should basically run the Service method of ModuleB once a day with the argument fed into ModuleA.Service.run.

The test I'd like to write:

import java.util.concurrent.atomic.AtomicInteger

import zio.clock.Clock
import zio.duration._
import zio.test.environment.TestClock
import zio.test.{DefaultRunnableSpec, assertCompletes, suite, testM}
import zio.{RIO, Task, ZIO}

object ExampleSpec extends DefaultRunnableSpec(ExampleSuite.suite1)

object ExampleSuite {

  val counter: AtomicInteger = new AtomicInteger(0)

  trait ModuleBTest extends ModuleB {
    override def moduleB: ModuleB.Service = new ModuleB.Service {
      override def run(arg: Int): RIO[Clock, Unit] = ZIO.effectTotal(counter.incrementAndGet())
    }
  }

  object ModuleATest extends ModuleALive with ModuleBTest

  def verifyExpectedInvocationCount(expectedInvocationCount: Int): Task[Unit] = {
    val actualInvocations = counter.get()
    if (counter.get() == expectedInvocationCount)
      ZIO.succeed(())
    else
      throw new Exception(s"expected invocation count: $expectedInvocationCount but was $actualInvocations")
  }

  val suite1 = suite("a")(
    testM("a should correctly schedule b") {
      for {
        _ <- ModuleATest.moduleA.schedule(42).fork
        _ <- TestClock.adjust(12 hours)
        _ <- verifyExpectedInvocationCount(1)
        _ <- TestClock.adjust(12 hours)
        _ <- verifyExpectedInvocationCount(2)
      } yield assertCompletes
    }
  )
}

I simplified the test using a counter, in reality I'd like to use mockito to verify the invocation count as well as the correct argument. However, this test does not work. In my understanding this is because of a race condition introduced by a timing overhead as described in https://zio.dev/docs/howto/howto_test_effects#testing-clock.

Now, there are examples of how to tackle this problem by using a Promise. I tried that by replacing the counter with a promise like so:

import java.util.concurrent.atomic.AtomicInteger

import zio.test.{DefaultRunnableSpec, assertCompletes, suite, testM}
import zio.{Promise, Task, UIO, ZIO}

object ExampleSpec extends DefaultRunnableSpec(ExampleSuite.suite1)

object ExampleSuite {

  val counter: AtomicInteger = new AtomicInteger(0)
  var promise: UIO[Promise[Unit, Int]] = Promise.make[Unit, Int]

  trait ModuleBTest extends ModuleB {
    override def moduleB: ModuleB.Service = new ModuleB.Service {
      override def run(arg: Int) = promise.map(_.succeed(counter.incrementAndGet))
    }
  }

  object ModuleATest extends ModuleALive with ModuleBTest

  def verifyExpectedInvocationCount(expectedInvocationCount: Int, actualInvocations: Int): Task[Unit] = {
    if (actualInvocations == expectedInvocationCount)
      ZIO.succeed(())
    else
      throw new Exception(s"expected invocation count: $expectedInvocationCount but was $actualInvocations")
  }

  val suite1 = suite("a")(
    testM("a should correctly schedule b") {
      for {
        _ <- ModuleATest.moduleA.schedule(42).fork
        p <- promise
        actualInvocationCount <- p.await
        _ <- verifyExpectedInvocationCount(expectedInvocationCount = 1, actualInvocationCount)
      } yield assertCompletes
    }
  )
}

Using this, the test won't terminate. However, I am pretty sure I am using the promise wrongly.

How would one approach this test scenario correctly?


Solution

  • In your example the type of promise is UIO[Promise[Unit, Int]] so you are creating a new promise each time. As a result, the promise your effect is completing is different than the one your test is waiting on, resulting in nontermination.

    To test this, you can do something like this:

    import zio.clock.Clock
    import zio.duration._
    import zio.test.environment.TestClock
    import zio.test.{ assertCompletes, suite, testM, DefaultRunnableSpec }
    import zio._
    
    object ExampleSpec extends DefaultRunnableSpec {
    
      trait ModuleA {
        def moduleA: ModuleA.Service
      }
    
      object ModuleA {
        trait Service {
          def schedule(arg: Int): RIO[Clock, Unit]
        }
      }
    
      trait ModuleALive extends ModuleA {
    
        def moduleB: ModuleB.Service
    
        override def moduleA: ModuleA.Service = new ModuleA.Service {
          override def schedule(arg: Int): RIO[Clock, Unit] =
            moduleB.run(arg).repeat(Schedule.spaced(1.day)).map(_ => ())
        }
      }
    
      trait ModuleB {
        def moduleB: ModuleB.Service
      }
    
      object ModuleB {
        trait Service {
          def run(arg: Int): RIO[Clock, Unit]
        }
      }
    
      trait ModuleBTest extends ModuleB {
        val counter: Ref[Int]
        val invocations: Queue[Int]
        override def moduleB: ModuleB.Service = new ModuleB.Service {
          override def run(arg: Int): UIO[Unit] =
            counter.updateAndGet(_ + 1).flatMap(invocations.offer).unit
        }
      }
    
      object ModuleATest {
        def apply(ref: Ref[Int], queue: Queue[Int]): ModuleALive with ModuleBTest =
          new ModuleALive with ModuleBTest {
            val counter     = ref
            val invocations = queue
          }
      }
    
      def verifyExpectedInvocationCount(invocations: Queue[Int], expected: Int): Task[Unit] =
        invocations.take.flatMap { actual =>
          if (actual == expected)
            ZIO.succeed(())
          else
            ZIO.fail(new Exception(s"expected invocation count: $expected but was $actual"))
        }
    
      def spec = suite("a")(
        testM("a should correctly schedule b") {
          for {
            counter     <- Ref.make(0)
            invocations <- Queue.unbounded[Int]
            moduleATest = ModuleATest(counter, invocations)
            _           <- moduleATest.moduleA.schedule(42).fork
            _           <- TestClock.adjust(12.hours)
            _           <- verifyExpectedInvocationCount(invocations, 1)
            _           <- TestClock.adjust(12.hours)
            _           <- verifyExpectedInvocationCount(invocations, 2)
          } yield assertCompletes
        }
      )
    }
    

    Since we want to wait for multiple effects to complete I use a Queue to coordinate them. A couple of other things to note: