ZIO Test: Add Gen Combinators

test/shared/src/main/scala/zio/test/Gen.scala

@@ -18,47 +18,99 @@ package zio.test
 
 import zio.{ UIO, ZIO }
 import zio.random._
-import zio.stream.ZStream
+import zio.stream.{ Stream, ZStream }
 
 /**
  * A `Gen[R, A]` represents a generator of values of type `A`, which requires
  * an environment `R`. Generators may be random or deterministic.
  */
 case class Gen[-R, +A](sample: ZStream[R, Nothing, Sample[R, A]]) { self =>
-  final def <*>[R1 <: R, B](that: Gen[R1, B]): Gen[R1, (A, B)] = self zip that
 
-  final def flatMap[R1 <: R, B](f: A => Gen[R1, B]): Gen[R1, B] =
-    Gen(sample.flatMap(sample => f(sample.value).sample))
+  final def <*>[R1 <: R, B](that: Gen[R1, B]): Gen[R1, (A, B)] =
+    self zip that
 
-  final def map[B](f: A => B): Gen[R, B] = Gen(sample.map(_.map(f)))
+  final def filter(f: A => Boolean): Gen[R, A] =
+    flatMap(a => if (f(a)) Gen.const(a) else Gen.empty)
 
-  final def zip[R1 <: R, B](that: Gen[R1, B]): Gen[R1, (A, B)] = self.flatMap(a => that.map(b => (a, b)))
+  final def flatMap[R1 <: R, B](f: A => Gen[R1, B]): Gen[R1, B] = Gen {
+    sample.flatMap { a =>
+      f(a.value).sample.map { b =>
+        val rest = a.shrink.flatMap(a => f(a).sample).flatMap(sample => ZStream(sample.value) ++ sample.shrink)
+        Sample(b.value, b.shrink ++ rest)
+      }
+    }
+  }
+
+  final def map[B](f: A => B): Gen[R, B] =
+    Gen(sample.map(_.map(f)))
+
+  final def noShrink: Gen[R, A] =
+    Gen(sample.map(sample => Sample.noShrink(sample.value)))
+
+  final def withShrink[R1 <: R, A1 >: A](f: A => ZStream[R1, Nothing, A1]): Gen[R1, A1] =
+    Gen(sample.map(sample => Sample(sample.value, f(sample.value) ++ sample.shrink.flatMap(f))))
+
+  final def zip[R1 <: R, B](that: Gen[R1, B]): Gen[R1, (A, B)] =
+    self.flatMap(a => that.map(b => (a, b)))
+
+  final def zipAll[R1 <: R, B](that: Gen[R1, B]): Gen[R1, (Option[A], Option[B])] =
+    Gen {
+      (self.sample.map(Some(_)) ++ Stream(None).forever)
+        .zip(that.sample.map(Some(_)) ++ Stream(None).forever)
+        .collectWhile {
+          case (Some(a), Some(b)) => a.zipWith(b) { case (a, b) => (Some(a), Some(b)) }
+          case (Some(a), None)    => a.map(a => (Some(a), None))
+          case (None, Some(b))    => b.map(b => (None, Some(b)))
+        }
+    }
 
   final def zipWith[R1 <: R, B, C](that: Gen[R1, B])(f: (A, B) => C): Gen[R1, C] =
     (self zip that) map f.tupled
 }
+
 object Gen {
 
+  /**
+   * A generator of booleans.
+   */
+  final val boolean: Gen[Random, Boolean] =
+    choose(false, true)
+
+  final def char(range: Range): Gen[Random, Char] =
+    integral[Char](range)
+
+  final def choose[A](a: A, as: A*): Gen[Random, A] =
+    oneOf(const(a), as.map(a => const(a)): _*)
+
   /**
    * A constant generator of the specified value.
    */
-  final def const[A](a: => A): Gen[Any, A] = Gen(ZStream.succeedLazy(Sample.noShrink(a)))
+  final def const[A](a: => A): Gen[Any, A] =
+    Gen(ZStream.succeedLazy(Sample.noShrink(a)))
 
   /**
    * A constant generator of the specified sample.
    */
-  final def constSample[R, A](sample: => Sample[R, A]): Gen[R, A] = fromEffectSample(ZIO.succeedLazy(sample))
+  final def constSample[R, A](sample: => Sample[R, A]): Gen[R, A] =
+    fromEffectSample(ZIO.succeedLazy(sample))
+
+  final def double(min: Double, max: Double): Gen[Random, Double] =
+    uniform.map(r => min + r * (max - min))
+
+  final val empty: Gen[Any, Nothing] =
+    Gen(Stream.empty)
 
   /**
    * Constructs a generator from an effect that constructs a value.
    */
   final def fromEffect[R, A](effect: ZIO[R, Nothing, A]): Gen[R, A] =
-    Gen(ZStream.fromEffect(effect.map(Sample.noShrink(_))))
+    Gen(ZStream.fromEffect(effect.map(Sample.noShrink)))
 
   /**
    * Constructs a generator from an effect that constructs a sample.
    */
-  final def fromEffectSample[R, A](effect: ZIO[R, Nothing, Sample[R, A]]): Gen[R, A] = Gen(ZStream.fromEffect(effect))
+  final def fromEffectSample[R, A](effect: ZIO[R, Nothing, Sample[R, A]]): Gen[R, A] =
+    Gen(ZStream.fromEffect(effect))
 
   /**
    * Constructs a deterministic generator that only generates the specified fixed values.
@@ -74,7 +126,7 @@ object Gen {
    * generator will not have any shrinking.
    */
   final def fromRandom[A](f: Random.Service[Any] => UIO[A]): Gen[Random, A] =
-    Gen(ZStream.fromEffect(ZIO.accessM[Random](r => f(r.random)).map(Sample.noShrink(_))))
+    Gen(ZStream.fromEffect(ZIO.accessM[Random](r => f(r.random)).map(Sample.noShrink)))
 
   /**
    * Constructs a generator from a function that uses randomness to produce a
@@ -84,26 +136,73 @@ object Gen {
     Gen(ZStream.fromEffect(ZIO.accessM[Random](r => f(r.random))))
 
   /**
-   * A generator of integral values inside the specified range: [start, end).
+   * A generator of integers inside the specified range: [start, end].
+   * The shrinker will shrink toward the lower end of the range ("smallest").
+   */
+  final def int(range: Range): Gen[Random, Int] =
+    integral[Int](range)
+
+  /**
+   * A generator of integral values inside the specified range: [start, end].
    * The shrinker will shrink toward the lower end of the range ("smallest").
    */
   final def integral[A](range: Range)(implicit I: Integral[A]): Gen[Random, A] =
     fromEffectSample(
-      nextInt(range.end - range.start)
+      nextInt(range.end - range.start + 1)
         .map(r => I.fromInt(r + range.start))
-        .map(int => Sample(int, ZStream.fromIterable(range.drop(1).reverse).map(I.fromInt(_))))
+        .map(int => Sample(int, ZStream.fromIterable(range.start.until(I.toInt(int)).reverse).map(n => I.fromInt(n))))
     )
 
+  final def listOf[R, A](range: Range)(gen: Gen[R, A]): Gen[R with Random, List[A]] =
+    int(range).flatMap(n => listOfN(n)(gen))
+
+  final def listOfN[R, A](n: Int)(gen: Gen[R, A]): Gen[R, List[A]] =
+    List.fill(n)(gen).foldRight[Gen[R, List[A]]](const(Nil))((a, gen) => a.zipWith(gen)(_ :: _))
+
+  final def optionOf[R, A](gen: Gen[R, A]): Gen[R with Random, Option[A]] =
+    oneOf(const(None), gen.map(Some(_)))
+
+  final def oneOf[R, A](a: Gen[R, A], as: Gen[R, A]*): Gen[R with Random, A] =
+    int(0 to as.length).flatMap(n => if (n == 0) a else as(n - 1))
+
   /**
    * A sized generator, whose size falls within the specified bounds.
    */
   final def sized[R <: Random, A](min: Int, max: Int)(f: Int => Gen[R, A]): Gen[R, A] =
     integral[Int](min to max).flatMap(f)
 
+  final def string[R](range: Range)(char: Gen[R, Char]): Gen[R with Random, String] =
+    listOf(range)(char).map(_.toString)
+
+  final def stringN[R](n: Int)(char: Gen[R, Char]): Gen[R, String] =
+    listOfN(n)(char).map(_.toString)
+
   /**
    * A generator of uniformly distributed doubles between [0, 1].
    *
    * TODO: Make Shrinker go toward `0`
    */
-  final def uniform: Gen[Random, Double] = Gen(ZStream.fromEffect(nextDouble.map(Sample.noShrink(_))))
+  final def uniform: Gen[Random, Double] =
+    Gen(ZStream.fromEffect(nextDouble.map(Sample.noShrink)))
+
+  final def union[R, A](gen1: Gen[R, A], gen2: Gen[R, A]): Gen[R, A] =
+    gen1.zipAll(gen2).flatMap {
+      case (a, a2) =>
+        Gen {
+          a.fold[Gen[R, A]](empty)(const(_)).sample ++
+            a2.fold[Gen[R, A]](empty)(const(_)).sample
+        }
+    }
+
+  /**
+   * A constant generator of the unit value.
+   */
+  final val unit: Gen[Any, Unit] =
+    const(())
+
+  final def vectorOf[R, A](range: Range)(gen: Gen[R, A]): Gen[R with Random, Vector[A]] =
+    listOf(range)(gen).map(_.toVector)
+
+  final def vectorOfN[R, A](n: Int)(gen: Gen[R, A]): Gen[R, Vector[A]] =
+    listOfN(n)(gen).map(_.toVector)
 }

test/shared/src/test/scala/zio/test/DefaultTestReporterSpec.scala

@@ -9,13 +9,13 @@ import zio.test.TestUtils.label
 object DefaultTestReporterSpec extends DefaultRuntime {
 
   def run(implicit ec: ExecutionContext): List[Future[(Boolean, String)]] = List(
-    label(reportSuccess, "DefaultTestReporter correctly reports a successful test"),
-    label(reportFailure, "DefaultTestReporter correctly reports a failed test"),
-    label(reportError, "DefaultTestReporter correctly reports an error in a test"),
-    label(reportSuite1, "DefaultTestReporter correctly reports successful test suite"),
-    label(reportSuite2, "DefaultTestReporter correctly reports failed test suite"),
-    label(reportSuites, "DefaultTestReporter correctly reports multiple test suites"),
-    label(simplePredicate, "DefaultTestReporter correctly reports failure of simple predicate")
+    label(reportSuccess, "correctly reports a successful test"),
+    label(reportFailure, "correctly reports a failed test"),
+    label(reportError, "correctly reports an error in a test"),
+    label(reportSuite1, "correctly reports successful test suite"),
+    label(reportSuite2, "correctly reports failed test suite"),
+    label(reportSuites, "correctly reports multiple test suites"),
+    label(simplePredicate, "correctly reports failure of simple predicate")
   )
 
   def makeTest[L](label: L)(assertion: => TestResult): ZSpec[Any, Nothing, L] =

test/shared/src/test/scala/zio/test/GenSpec.scala

@@ -0,0 +1,163 @@
+package zio.test
+
+import scala.concurrent.{ ExecutionContext, Future }
+
+import zio.{ DefaultRuntime, Managed, UIO, ZIO }
+import zio.random.Random
+import zio.stream.ZStream
+import zio.test.mock.MockRandom
+import zio.test.TestUtils.label
+
+object GenSpec extends DefaultRuntime {
+
+  def run(implicit ec: ExecutionContext): List[Future[(Boolean, String)]] = List(
+    label(monadLeftIdentity, "monad left identity"),
+    label(monadRightIdentity, "monad right identity"),
+    label(monadAssociativity, "monad associativity"),
+    label(booleanGeneratesTrueAndFalse, "boolean generates true and false"),
+    label(booleanShrinksToFalse, "boolean shrinks to false"),
+    label(charGeneratesValuesInRange, "char generates values in range"),
+    label(constGeneratesConstantValue, "const generates constant value"),
+    label(doubleGeneratesValuesInRange, "double generates values in range"),
+    label(filterFiltersValuesAccordingToPredicate, "filter filters values according to predicate"),
+    label(fromIterableGeneratorsCompose, "finite generators compose"),
+    label(intGeneratesValuesInRange, "int generates values in range"),
+    label(intShrinksToBottomOfRange, "int shrinks to bottom of range"),
+    label(listOfShrinksToSmallestLength, "listOf shrinks to smallest length"),
+    label(noShrinkRemovesShrinkingLogic, "no shrink removes shrinking logic"),
+    label(uniformGeneratesValuesInRange, "uniform generates values between 0 and 1"),
+    label(unionCombinesTwoGenerators, "union combines two generators"),
+    label(unit, "unit generates the constant unit value"),
+    label(withShrinkAddsShrinkingLogic, "with shrink adds shrinking logic"),
+    label(zipShrinksCorrectly, "zip shrinks correctly"),
+    label(zipWithShrinksCorrectly, "zipWith shrinks correctly")
+  )
+
+  val smallInt = Gen.int(-10 to 10)
+
+  def monadLeftIdentity: Future[Boolean] =
+    checkEqual(smallInt.flatMap(a => Gen.const(a)), smallInt)
+
+  def monadRightIdentity: Future[Boolean] = {
+    val n                           = 10
+    def f(n: Int): Gen[Random, Int] = Gen.int(-n to n)
+    checkEqual(Gen.const(n).flatMap(f), f(n))
+  }
+
+  def monadAssociativity: Future[Boolean] = {
+    val fa = Gen.int(0 to 2)
+    def f(p: Int): Gen[Random, (Int, Int)] =
+      Gen.const(p) <*> Gen.int(0 to 3)
+    def g(p: (Int, Int)): Gen[Random, (Int, Int, Int)] =
+      Gen.const(p).zipWith(Gen.int(0 to 5)) { case ((x, y), z) => (x, y, z) }
+    checkEqual(fa.flatMap(f).flatMap(g), fa.flatMap(a => f(a).flatMap(g)))
+  }
+
+  def booleanGeneratesTrueAndFalse: Future[Boolean] =
+    checkSample(Gen.boolean)(ps => ps.exists(identity) && ps.exists(!_))
+
+  def booleanShrinksToFalse: Future[Boolean] =
+    checkShrink(Gen.boolean)(false)
+
+  def charGeneratesValuesInRange: Future[Boolean] =
+    checkSample(Gen.char(33 to 123))(_.forall(n => 33 <= n && n <= 123))
+
+  def constGeneratesConstantValue: Future[Boolean] =
+    checkSample(Gen.const("constant"))(_.forall(_ == "constant"))
+
+  def doubleGeneratesValuesInRange: Future[Boolean] =
+    checkSample(Gen.double(5.0, 9.0))(_.forall(n => 5.0 <= n && n < 9.0))
+
+  def filterFiltersValuesAccordingToPredicate: Future[Boolean] =
+    checkSample(smallInt.filter(_ % 2 == 0))(_.forall(_ % 2 == 0))
+
+  def fromIterableGeneratorsCompose: Future[Boolean] = {
+    val exhaustive = Gen.fromIterable(1 to 6)
+    val actual     = exhaustive.zipWith(exhaustive)(_ + _)
+    val expected   = (1 to 6).flatMap(x => (1 to 6).map(y => x + y))
+    checkFinite(actual)(_ == expected)
+  }
+
+  def intGeneratesValuesInRange: Future[Boolean] =
+    checkSample(smallInt)(_.forall(n => -10 <= n && n <= 10))
+
+  def intShrinksToBottomOfRange: Future[Boolean] =
+    checkShrink(smallInt)(-10)
+
+  def listOfShrinksToSmallestLength: Future[Boolean] =
+    checkShrink(Gen.listOf(0 to 10)(Gen.boolean))(Nil)
+
+  def noShrinkRemovesShrinkingLogic: Future[Boolean] =
+    forAll(smallInt.noShrink.sample.flatMap(_.shrink).runCollect.map(_.isEmpty))
+
+  def uniformGeneratesValuesInRange: Future[Boolean] =
+    checkSample(Gen.uniform)(_.forall(n => 0.0 <= n && n < 1.0))
+
+  def unionCombinesTwoGenerators: Future[Boolean] = {
+    val as = Gen.fromIterable(2 to 3)
+    val bs = Gen.fromIterable(5 to 7)
+    val cs = Gen.union(as, bs)
+    checkFinite(cs)(_ == List(2, 5, 3, 6, 7))
+  }
+
+  def unit: Future[Boolean] =
+    checkSample(Gen.unit)(_.forall(_ => true))
+
+  def withShrinkAddsShrinkingLogic: Future[Boolean] = {
+    def shrink(bound: Int): Int => ZStream[Any, Nothing, Int] =
+      n => if (n > bound) ZStream(n - 1) ++ shrink(bound)(n - 1) else ZStream.empty
+    checkShrink(Gen.const(10).withShrink(shrink(0)))(0)
+  }
+
+  def zipShrinksCorrectly: Future[Boolean] =
+    checkShrink(smallInt <*> smallInt)((-10, -10))
+
+  def zipWithShrinksCorrectly: Future[Boolean] =
+    checkShrink(smallInt.zipWith(smallInt)(_ + _))(-20)
+
+  def checkEqual[A](left: Gen[Random, A], right: Gen[Random, A]): Future[Boolean] =
+    unsafeRunToFuture(equal(left, right))
+
+  def checkSample[A](gen: Gen[Random, A])(f: List[A] => Boolean): Future[Boolean] =
+    unsafeRunToFuture(sample(gen).map(f))
+
+  def checkFinite[A](gen: Gen[Random, A])(f: List[A] => Boolean): Future[Boolean] =
+    unsafeRunToFuture(gen.sample.map(_.value).runCollect.map(f))
+
+  def checkShrink[A](gen: Gen[Random, A])(a: A): Future[Boolean] =
+    unsafeRunToFuture(alwaysShrinksTo(gen)(a: A))
+
+  def sample[A](gen: Gen[Random, A]): ZIO[Random, Nothing, List[A]] =
+    gen.sample.map(_.value).forever.take(100).runCollect
+
+  def alwaysShrinksTo[A](gen: Gen[Random, A])(a: A): ZIO[Random, Nothing, Boolean] =
+    ZIO.collectAll(List.fill(100)(shrinksTo(gen))).map(_.forall(_ == a))
+
+  def shrinksTo[A](gen: Gen[Random, A]): ZIO[Random, Nothing, A] =
+    shrinks(gen).map(_.reverse.head)
+
+  def shrinks[A](gen: Gen[Random, A]): ZIO[Random, Nothing, List[A]] =
+    gen.sample.flatMap(sample => ZStream(sample.value) ++ sample.shrink).take(1000).runCollect
+
+  def equal[A](left: Gen[Random, A], right: Gen[Random, A]): UIO[Boolean] =
+    equalSample(left, right).zipWith(equalShrink(left, right))(_ && _)
+
+  def forAll[E <: Throwable, A](zio: ZIO[Random, E, Boolean]): Future[Boolean] =
+    unsafeRunToFuture(ZIO.collectAll(List.fill(100)(zio)).map(_.forall(identity)))
+
+  def equalSample[A](left: Gen[Random, A], right: Gen[Random, A]): UIO[Boolean] = {
+    val mockRandom = Managed.fromEffect(MockRandom.make(MockRandom.DefaultData))
+    for {
+      leftSample  <- sample(left).provideManaged(mockRandom)
+      rightSample <- sample(right).provideManaged(mockRandom)
+    } yield leftSample == rightSample
+  }
+
+  def equalShrink[A](left: Gen[Random, A], right: Gen[Random, A]): UIO[Boolean] = {
+    val mockRandom = Managed.fromEffect(MockRandom.make(MockRandom.DefaultData))
+    for {
+      leftShrinks  <- ZIO.collectAll(List.fill(100)(shrinks(left))).provideManaged(mockRandom)
+      rightShrinks <- ZIO.collectAll(List.fill(100)(shrinks(right))).provideManaged(mockRandom)
+    } yield leftShrinks == rightShrinks
+  }
+}