1 files changed, 168 insertions, 10 deletions
diff --git a/src/test/java/com/squareup/javapoet/MethodSpecTest.java b/src/test/java/com/squareup/javapoet/MethodSpecTest.java
index 5dfabaa..b768351 100644
--- a/src/test/java/com/squareup/javapoet/MethodSpecTest.java
+++ b/src/test/java/com/squareup/javapoet/MethodSpecTest.java
@@ -21,8 +21,9 @@ import java.io.IOException;
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Target;
 import java.util.Arrays;
-import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeoutException;
 import javax.lang.model.element.ExecutableElement;
@@ -37,6 +38,8 @@ import org.junit.Test;
 
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
+import static com.squareup.javapoet.MethodSpec.CONSTRUCTOR;
+import static com.squareup.javapoet.TestUtil.findFirst;
 import static javax.lang.model.util.ElementFilter.methodsIn;
 import static org.junit.Assert.fail;
 
@@ -55,15 +58,6 @@ public final class MethodSpecTest {
     return elements.getTypeElement(clazz.getCanonicalName());
   }
 
-  private ExecutableElement findFirst(Collection<ExecutableElement> elements, String name) {
-    for (ExecutableElement executableElement : elements) {
-      if (executableElement.getSimpleName().toString().equals(name)) {
-        return executableElement;
-      }
-    }
-    throw new IllegalArgumentException(name + " not found in " + elements);
-  }
-
   @Test public void nullAnnotationsAddition() {
     try {
       MethodSpec.methodBuilder("doSomething").addAnnotations(null);
@@ -270,6 +264,59 @@ public final class MethodSpecTest {
     assertThat(a.hashCode()).isEqualTo(b.hashCode());
   }
 
+  @Test public void withoutParameterJavaDoc() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("getTaco")
+        .addModifiers(Modifier.PRIVATE)
+        .addParameter(TypeName.DOUBLE, "money")
+        .addJavadoc("Gets the best Taco\n")
+        .build();
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "/**\n"
+        + " * Gets the best Taco\n"
+        + " */\n"
+        + "private void getTaco(double money) {\n"
+        + "}\n");
+  }
+
+  @Test public void withParameterJavaDoc() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("getTaco")
+        .addParameter(ParameterSpec.builder(TypeName.DOUBLE, "money")
+            .addJavadoc("the amount required to buy the taco.\n")
+            .build())
+        .addParameter(ParameterSpec.builder(TypeName.INT, "count")
+            .addJavadoc("the number of Tacos to buy.\n")
+            .build())
+        .addJavadoc("Gets the best Taco money can buy.\n")
+        .build();
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "/**\n"
+        + " * Gets the best Taco money can buy.\n"
+        + " *\n"
+        + " * @param money the amount required to buy the taco.\n"
+        + " * @param count the number of Tacos to buy.\n"
+        + " */\n"
+        + "void getTaco(double money, int count) {\n"
+        + "}\n");
+  }
+
+  @Test public void withParameterJavaDocAndWithoutMethodJavadoc() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("getTaco")
+        .addParameter(ParameterSpec.builder(TypeName.DOUBLE, "money")
+            .addJavadoc("the amount required to buy the taco.\n")
+            .build())
+        .addParameter(ParameterSpec.builder(TypeName.INT, "count")
+            .addJavadoc("the number of Tacos to buy.\n")
+            .build())
+        .build();
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "/**\n"
+        + " * @param money the amount required to buy the taco.\n"
+        + " * @param count the number of Tacos to buy.\n"
+        + " */\n"
+        + "void getTaco(double money, int count) {\n"
+        + "}\n");
+  }
+
   @Test public void duplicateExceptionsIgnored() {
     ClassName ioException = ClassName.get(IOException.class);
     ClassName timeoutException = ClassName.get(TimeoutException.class);
@@ -302,4 +349,115 @@ public final class MethodSpecTest {
       assertThat(e.getMessage()).isEqualTo("modifiers == null");
     }
   }
+
+  @Test public void modifyMethodName() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("initialMethod")
+        .build()
+        .toBuilder()
+        .setName("revisedMethod")
+        .build();
+
+    assertThat(methodSpec.toString()).isEqualTo("" + "void revisedMethod() {\n" + "}\n");
+  }
+
+  @Test public void modifyAnnotations() {
+    MethodSpec.Builder builder = MethodSpec.methodBuilder("foo")
+            .addAnnotation(Override.class)
+            .addAnnotation(SuppressWarnings.class);
+
+    builder.annotations.remove(1);
+    assertThat(builder.build().annotations).hasSize(1);
+  }
+
+  @Test public void modifyModifiers() {
+    MethodSpec.Builder builder = MethodSpec.methodBuilder("foo")
+            .addModifiers(Modifier.PUBLIC, Modifier.STATIC);
+
+    builder.modifiers.remove(1);
+    assertThat(builder.build().modifiers).containsExactly(Modifier.PUBLIC);
+  }
+
+  @Test public void modifyParameters() {
+    MethodSpec.Builder builder = MethodSpec.methodBuilder("foo")
+            .addParameter(int.class, "source");
+
+    builder.parameters.remove(0);
+    assertThat(builder.build().parameters).isEmpty();
+  }
+
+  @Test public void modifyTypeVariables() {
+    TypeVariableName t = TypeVariableName.get("T");
+    MethodSpec.Builder builder = MethodSpec.methodBuilder("foo")
+            .addTypeVariable(t)
+            .addTypeVariable(TypeVariableName.get("V"));
+
+    builder.typeVariables.remove(1);
+    assertThat(builder.build().typeVariables).containsExactly(t);
+  }
+
+  @Test public void ensureTrailingNewline() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("method")
+        .addCode("codeWithNoNewline();")
+        .build();
+
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "void method() {\n"
+        + "  codeWithNoNewline();\n"
+        + "}\n");
+  }
+
+  /** Ensures that we don't add a duplicate newline if one is already present. */
+  @Test public void ensureTrailingNewlineWithExistingNewline() {
+    MethodSpec methodSpec = MethodSpec.methodBuilder("method")
+        .addCode("codeWithNoNewline();\n") // Have a newline already, so ensure we're not adding one
+        .build();
+
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "void method() {\n"
+        + "  codeWithNoNewline();\n"
+        + "}\n");
+  }
+
+  @Test public void controlFlowWithNamedCodeBlocks() {
+    Map<String, Object> m = new HashMap<>();
+    m.put("field", "valueField");
+    m.put("threshold", "5");
+
+    MethodSpec methodSpec = MethodSpec.methodBuilder("method")
+        .beginControlFlow(named("if ($field:N > $threshold:L)", m))
+        .nextControlFlow(named("else if ($field:N == $threshold:L)", m))
+        .endControlFlow()
+        .build();
+
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "void method() {\n"
+        + "  if (valueField > 5) {\n"
+        + "  } else if (valueField == 5) {\n"
+        + "  }\n"
+        + "}\n");
+  }
+
+  @Test public void doWhileWithNamedCodeBlocks() {
+    Map<String, Object> m = new HashMap<>();
+    m.put("field", "valueField");
+    m.put("threshold", "5");
+
+    MethodSpec methodSpec = MethodSpec.methodBuilder("method")
+        .beginControlFlow("do")
+        .addStatement(named("$field:N--", m))
+        .endControlFlow(named("while ($field:N > $threshold:L)", m))
+        .build();
+
+    assertThat(methodSpec.toString()).isEqualTo(""
+        + "void method() {\n" +
+        "  do {\n" +
+        "    valueField--;\n" +
+        "  } while (valueField > 5);\n" +
+        "}\n");
+  }
+
+  private static CodeBlock named(String format, Map<String, ?> args){
+    return CodeBlock.builder().addNamed(format, args).build();
+  }
+
 }