Add more tests for LambdaTypeName (#478)

* Add more tests for LambdaTypeName * Wrap LambdaTypeName return type in brackets
square · Oct 7, 2018 · c348f8e · c348f8e
1 parent 6921368
commit c348f8e
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt b/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt
@@ -70,7 +70,7 @@ class LambdaTypeName internal constructor(
     }
 
     parameters.emit(out)
-    out.emitCode(" -> %T", returnType)
+    out.emitCode(if (returnType is LambdaTypeName) " -> (%T)" else " -> %T", returnType)
 
     if (nullable) {
       out.emit(")")

diff --git a/src/test/java/com/squareup/kotlinpoet/LambdaTypeNameTest.kt b/src/test/java/com/squareup/kotlinpoet/LambdaTypeNameTest.kt
@@ -76,4 +76,26 @@ class LambdaTypeNameTest {
           returnType = Unit::class.asTypeName())
     }.hasMessageThat().isEqualTo("Parameters with default values are not allowed")
   }
+
+  @Test fun lambdaReturnType() {
+    val returnTypeName = LambdaTypeName.get(
+        parameters = *arrayOf(Int::class.asTypeName()),
+        returnType = Unit::class.asTypeName())
+    val typeName = LambdaTypeName.get(
+        parameters = *arrayOf(Int::class.asTypeName()),
+        returnType = returnTypeName)
+    assertThat(typeName.toString())
+        .isEqualTo("(kotlin.Int) -> ((kotlin.Int) -> kotlin.Unit)")
+  }
+
+  @Test fun lambdaParameterType() {
+    val parameterTypeName = LambdaTypeName.get(
+        parameters = *arrayOf(Int::class.asTypeName()),
+        returnType = Int::class.asTypeName())
+    val typeName = LambdaTypeName.get(
+        parameters = *arrayOf(parameterTypeName),
+        returnType = Unit::class.asTypeName())
+    assertThat(typeName.toString())
+        .isEqualTo("((kotlin.Int) -> kotlin.Int) -> kotlin.Unit")
+  }
 }