zio · kyri-petrou · Mar 28, 2025 · Mar 24, 2025 · Mar 24, 2025 · Mar 24, 2025
diff --git a/test-magnolia-tests/src/test/scala-3/zio/test/magnolia/DeriveGenSpec_3.scala b/test-magnolia-tests/src/test/scala-3/zio/test/magnolia/DeriveGenSpec_3.scala
@@ -0,0 +1,67 @@
+package zio.test.magnolia
+
+import zio._
+import zio.test.Assertion._
+import zio.test.GenUtils._
+import zio.test.magnolia.DeriveGen._
+import zio.test._
+import zio.test.TestAspect.samples
+
+import java.time.{Instant, LocalDate, LocalDateTime, LocalTime}
+import java.util.UUID
+
+object DeriveGenSpec_3 extends ZIOBaseSpec {
+
+  sealed trait Color
+  object Color {
+    case object Red   extends Color
+    case object Green extends Color
+    case object Blue  extends Color
+  }
+
+  val genColor: Gen[Any, Color] = DeriveGen[Color]
+
+  type UnionType = Color | Boolean
+  given DeriveGen[UnionType]            = DeriveGen.unionType[UnionType]
+  val anyUnionType: Gen[Any, UnionType] = DeriveGen[UnionType]
+
+  type NestedUnionType = String | UnionType
+  given DeriveGen[NestedUnionType]                  = DeriveGen.unionType[NestedUnionType]
+  val anyNestedUnionType: Gen[Any, NestedUnionType] = DeriveGen[NestedUnionType]
+
+  def assertDeriveGen[A](implicit ev: DeriveGen[A]): TestResult = assertCompletes
+
+  def spec = suite("DeriveGenSpec_3")(
+    suite("union type")(
+      test("derivation")(assertDeriveGen[UnionType]),
+      test("generates Color values") {
+        check(Gen.listOfN(100)(anyUnionType)) { vs =>
+          assertTrue(vs.exists(_.isInstanceOf[Color]))
+        }
+      },
+      test("generates Boolean values") {
+        check(Gen.listOfN(100)(anyUnionType)) { vs =>
+          assertTrue(vs.exists(_.isInstanceOf[Boolean]))
+        }
+      }
+    ),
+    suite("nested union type")(
+      test("derivation")(assertDeriveGen[NestedUnionType]),
+      test("generates Color values") {
+        check(Gen.listOfN(100)(anyNestedUnionType)) { vs =>
+          assertTrue(vs.exists(_.isInstanceOf[Color]))
+        }
+      },
+      test("generates Boolean values") {
+        check(Gen.listOfN(100)(anyNestedUnionType)) { vs =>
+          assertTrue(vs.exists(_.isInstanceOf[Boolean]))
+        }
+      },
+      test("generates String values") {
+        check(Gen.listOfN(100)(anyNestedUnionType)) { vs =>
+          assertTrue(vs.exists(_.isInstanceOf[String]))
+        }
+      }
+    )
+  )
+}
diff --git a/test-magnolia/src/main/scala-3/zio/test/magnolia/DeriveGen.scala b/test-magnolia/src/main/scala-3/zio/test/magnolia/DeriveGen.scala
@@ -38,6 +38,18 @@ object DeriveGen {
       val derive: Gen[Any, A] = gen
     }
 
+  /**
+   * Util to derive a `DeriveGen` instance for a union type
+   *
+   * Usage example:
+   * {{{
+   *  type StringOrInt = String | Int
+   *  given DeriveGen[StringOrInt] = DeriveGen.unionType[StringOrInt]
+   *  lazy val anyStringOrInt: Gen[Any, StringOrInt] = DeriveGen[UnionType]
+   * }}}
+   */
+  inline def unionType[T]: DeriveGen[T] = ${ TypeUnionDerivation.typeUnionDeriveGen[T] }
+
   given DeriveGen[Boolean]       = instance(Gen.boolean)
   given DeriveGen[Byte]          = instance(Gen.byte)
   given DeriveGen[Char]          = instance(Gen.char)

diff --git a/test-magnolia/src/main/scala-3/zio/test/magnolia/TypeUnionDerivation.scala b/test-magnolia/src/main/scala-3/zio/test/magnolia/TypeUnionDerivation.scala
@@ -0,0 +1,65 @@
+package zio.test.magnolia
+
+import scala.quoted.*
+
+import zio.test.Gen
+
+/**
+ * Implementation copied/adapted from Caliban
+ *
+ * See:
+ *   - https://github.com/ghostdogpr/caliban/pull/2215
+ *   - https://github.com/ghostdogpr/caliban/blob/v2.10.0/core/src/main/scala-3/caliban/schema/TypeUnionDerivation.scala
+ */
+object TypeUnionDerivation {
+  inline def derived[T]: DeriveGen[T] = ${ typeUnionDeriveGen[T] }
+
+  def typeUnionDeriveGen[T](using Quotes, Type[T]): Expr[DeriveGen[T]] = {
+    import quotes.reflect.*
+
+    val typeName = TypeRepr.of[T].show
+
+    if (typeName.contains("|")) {
+      report.error(
+        s"You must explicitly add type parameter to derive DeriveGen for a union type in order to capture the name of the type alias"
+      )
+    }
+
+    class TypeAndDeriveGen[A](val deriveGen: Expr[DeriveGen[A]], val tpe: Type[A])
+
+    def rec[A](using tpe: Type[A]): List[TypeAndDeriveGen[?]] =
+      TypeRepr.of(using tpe).dealias match {
+        case OrType(l, r) =>
+          rec(using l.asType.asInstanceOf[Type[Any]]) ++ rec(using r.asType.asInstanceOf[Type[Any]])
+        case otherRepr =>
+          val expr: TypeAndDeriveGen[A] =
+            Expr.summon[DeriveGen[A]] match {
+              case Some(foundDeriveGen) =>
+                TypeAndDeriveGen[A](foundDeriveGen, otherRepr.asType.asInstanceOf[Type[A]])
+              case None =>
+                val otherString: String = otherRepr.show
+                quotes.reflect.report.errorAndAbort(s"Couldn't resolve DeriveGen[$otherString]")
+            }
+
+          List(expr)
+      }
+
+    val typeAndDeriveGens: List[TypeAndDeriveGen[?]] = rec[T]
+
+    val deriveGenByTypeNameList: Expr[List[DeriveGen[Any]]] = Expr.ofList(
+      typeAndDeriveGens.map { case (t: TypeAndDeriveGen[a]) =>
+        given Type[a] = t.tpe
+        '{ ${ t.deriveGen }.asInstanceOf[DeriveGen[Any]] }
+      }
+    )
+
+    '{
+      new DeriveGen[T] {
+        private lazy val gen: Gen[Any, T] =
+          Gen.oneOf[Any, T](${ deriveGenByTypeNameList }.map(_.derive).asInstanceOf[List[Gen[Any, T]]]*)
+
+        def derive: Gen[Any, T] = gen
+      }
+    }
+  }
+}