Fix primitive detection for mock generator

apicollective · Jun 21, 2024 · b326c1c · b326c1c
1 parent e16a5fe
commit b326c1c
Showing 1 changed file with 18 additions and 11 deletions.
diff --git a/scala-generator/src/main/scala/models/generator/mock/MockFactoriesGenerator.scala b/scala-generator/src/main/scala/models/generator/mock/MockFactoriesGenerator.scala
@@ -54,23 +54,30 @@ object MockFactoriesGenerator extends CodeGenerator {
   }
 
   private[mock] def makeUnion(union: ScalaUnion): String = {
-    val impl = union.types.headOption match {
-      case Some(typ) => {
-        val value = mockValue(typ.datatype)
-        typ.datatype match {
-          case p: ScalaPrimitive => {
-            // Union type wrapping a primitive
-            val className = Text.pascalCase(union.name) + Text.pascalCase(p.shortName)
-            union.ssd.namespaces.unions + "." + className + s"($value)"
-          }
-          case _ => value
+    val impl = union.types.headOption.map(_.datatype) match {
+      case Some(datatype: ScalaPrimitive) => {
+        val value = mockValue(datatype)
+        if (needsWrapper(datatype)) {
+          val className = Text.pascalCase(union.name) + Text.pascalCase(datatype.shortName)
+          union.ssd.namespaces.unions + "." + className + s"($value)"
+        } else {
+          value
         }
       }
-      case None => "???"
+      case _ => "???" // no types defined
     }
     s"def make${union.name}(): ${union.qualifiedName} = $impl"
   }
 
+  private[this] def needsWrapper(datatype: ScalaDatatype): Boolean = {
+    datatype match {
+      case _: ScalaPrimitive.Model | _: ScalaPrimitive.GeneratedModel | _: ScalaPrimitive.Union => {
+        false
+      }
+      case _ => true
+    }
+  }
+
   private[mock] def mockValue(
     datatype: ScalaDatatype,
     limitation: Option[ScalaField.Limitation] = None,