Fix bug where only ClassTrees with kind CLASS are considered by TypeScanner, resulting in a broken generatesSources() for interfaces, enums, and annotation types.

cgruber · cgruber · commit dd69f2a8553e · 2014-07-17T13:45:41.000-07:00
------------- Created by MOE: http://code.google.com/p/moe-java MOE_MIGRATED_REVID=65448993
diff --git a/src/main/java/com/google/testing/compile/TypeScanner.java b/src/main/java/com/google/testing/compile/TypeScanner.java
@@ -15,10 +15,13 @@
  */
 package com.google.testing.compile;
 
+import static com.sun.source.tree.Tree.Kind.ANNOTATION_TYPE;
 import static com.sun.source.tree.Tree.Kind.CLASS;
+import static com.sun.source.tree.Tree.Kind.ENUM;
 import static com.sun.source.tree.Tree.Kind.COMPILATION_UNIT;
 import static com.sun.source.tree.Tree.Kind.EXPRESSION_STATEMENT;
 import static com.sun.source.tree.Tree.Kind.IDENTIFIER;
+import static com.sun.source.tree.Tree.Kind.INTERFACE;
 import static com.sun.source.tree.Tree.Kind.MEMBER_SELECT;
 
 import com.google.common.base.Function;
@@ -59,7 +62,8 @@ static ImmutableSet<String> getTopLevelTypes(CompilationUnitTree t) {
   static final class NameVisitor extends TreeScanner<Set<String>, Void> {
 
     private static final Set<Tree.Kind> RELEVANT_KINDS = Sets.immutableEnumSet(
-        CLASS, COMPILATION_UNIT, EXPRESSION_STATEMENT, IDENTIFIER, MEMBER_SELECT);
+        ANNOTATION_TYPE, CLASS, ENUM, COMPILATION_UNIT, EXPRESSION_STATEMENT, IDENTIFIER, INTERFACE,
+        MEMBER_SELECT);
 
     @Override
     public Set<String> scan(Tree node, Void v) {
diff --git a/src/test/java/com/google/testing/compile/TypeScannerTest.java b/src/test/java/com/google/testing/compile/TypeScannerTest.java
@@ -108,6 +108,47 @@ public void getTopLevelTypes_manySimpleTypes() {
         "HelloWorld", "HelperWorld");
   }
 
+  @Test
+  public void getTopLevelTypes_worksForAnnotationTypes() {
+    CompilationUnitTree compilation = compileLines("test.HelloWorld",
+        "package path.to.test;",
+        "import java.util.List;",
+        "",
+        "public @interface HelloWorld {}");
+
+    ASSERT.that(TypeScanner.getTopLevelTypes(compilation)).has().exactly(
+        "path.to.test.HelloWorld");
+  }
+
+  @Test
+  public void getTopLevelTypes_worksForEnums() {
+    CompilationUnitTree compilation = compileLines("test.HelloWorld",
+        "package path.to.test;",
+        "import java.util.List;",
+        "",
+        "public enum HelloWorld {",
+        "  HELLO,",
+        "  WORLD;",
+        "}");
+
+    ASSERT.that(TypeScanner.getTopLevelTypes(compilation)).has().exactly(
+        "path.to.test.HelloWorld");
+  }
+
+  @Test
+  public void getTopLevelTypes_worksForInterfaces() {
+    CompilationUnitTree compilation = compileLines("test.HelloWorld",
+        "package path.to.test;",
+        "import java.util.List;",
+        "",
+        "public interface HelloWorld {",
+        "  public String getSalutation();",
+        "}");
+
+    ASSERT.that(TypeScanner.getTopLevelTypes(compilation)).has().exactly(
+        "path.to.test.HelloWorld");
+  }
+
   @Test
   public void getTopLevelTypes_worksForNull() {
     ASSERT.that(TypeScanner.getTopLevelTypes(null)).isEmpty();
