Look for annotations on methods as well as classes

Currently using something like `@FlakyTest` only works if a whole class is annotated.

It can be preferable to just annotate the single flaky method, rather than the whole class. This allows that.

java/gazelle/generate.go

@@ -415,7 +415,13 @@ func addFilteringOutOwnPackage(to *sorted_set.SortedSet[types.PackageName], from
 
 func accumulateJavaFile(cfg *javaconfig.Config, testJavaFiles, testHelperJavaFiles *sorted_set.SortedSet[javaFile], separateTestJavaFiles map[javaFile]separateJavaTestReasons, file javaFile, perClassMetadata map[string]java.PerClassMetadata, log zerolog.Logger) {
 	if cfg.IsJavaTestFile(filepath.Base(file.pathRelativeToBazelWorkspaceRoot)) {
-		annotationClassNames := perClassMetadata[file.ClassName().FullyQualifiedClassName()].AnnotationClassNames
+		annotationClassNames := sorted_set.NewSortedSet[string](nil)
+		metadataForClass := perClassMetadata[file.ClassName().FullyQualifiedClassName()]
+		annotationClassNames.AddAll(metadataForClass.AnnotationClassNames)
+		for _, key := range metadataForClass.MethodAnnotationClassNames.Keys() {
+			annotationClassNames.AddAll(metadataForClass.MethodAnnotationClassNames.Values(key))
+		}
+
 		perFileAttrs := make(map[string]bzl.Expr)
 		wrapper := ""
 		for _, annotationClassName := range annotationClassNames.SortedSlice() {

java/gazelle/lang.go

@@ -177,7 +177,7 @@ func (l javaLang) Loads() []rule.LoadInfo {
 	for _, name := range s.Keys() {
 		loads = append(loads, rule.LoadInfo{
 			Name:    name,
-			Symbols: s.Values(name),
+			Symbols: s.SortedValues(name),
 		})
 	}
 	return loads

java/gazelle/private/java/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
     # Allow visibility for plugins like Kotlin that don't live in this repo
     visibility = ["//visibility:public"],
     deps = [
+        "//java/gazelle/private/sorted_multiset",
         "//java/gazelle/private/sorted_set",
         "//java/gazelle/private/types",
     ],

java/gazelle/private/java/package.go

@@ -1,6 +1,7 @@
 package java
 
 import (
+	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/sorted_multiset"
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/sorted_set"
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/types"
 )
@@ -21,5 +22,6 @@ type Package struct {
 }
 
 type PerClassMetadata struct {
-	AnnotationClassNames *sorted_set.SortedSet[string]
+	AnnotationClassNames       *sorted_set.SortedSet[string]
+	MethodAnnotationClassNames *sorted_multiset.SortedMultiSet[string, string]
 }

java/gazelle/private/javaparser/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
         "//java/gazelle/private/java",
         "//java/gazelle/private/javaparser/proto/gazelle/java/javaparser/v0:javaparser",
         "//java/gazelle/private/servermanager",
+        "//java/gazelle/private/sorted_multiset",
         "//java/gazelle/private/sorted_set",
         "//java/gazelle/private/types",
         "@com_github_rs_zerolog//:zerolog",

java/gazelle/private/javaparser/javaparser.go

@@ -8,6 +8,7 @@ import (
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/java"
 	pb "github.com/bazel-contrib/rules_jvm/java/gazelle/private/javaparser/proto/gazelle/java/javaparser/v0"
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/servermanager"
+	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/sorted_multiset"
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/sorted_set"
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/types"
 	"github.com/rs/zerolog"
@@ -64,8 +65,15 @@ func (r Runner) ParsePackage(ctx context.Context, in *ParsePackageRequest) (*jav
 
 	perClassMetadata := make(map[string]java.PerClassMetadata, len(resp.GetPerClassMetadata()))
 	for k, v := range resp.GetPerClassMetadata() {
+		methodAnnotationClassNames := sorted_multiset.NewSortedMultiSet[string, string]()
+		for method, perMethod := range v.GetPerMethodMetadata() {
+			for _, annotation := range perMethod.AnnotationClassNames {
+				methodAnnotationClassNames.Add(method, annotation)
+			}
+		}
 		metadata := java.PerClassMetadata{
-			AnnotationClassNames: sorted_set.NewSortedSet(v.GetAnnotationClassNames()),
+			AnnotationClassNames:       sorted_set.NewSortedSet(v.GetAnnotationClassNames()),
+			MethodAnnotationClassNames: methodAnnotationClassNames,
 		}
 		perClassMetadata[k] = metadata
 	}

java/gazelle/private/javaparser/proto/gazelle/java/javaparser/v0/javaparser.proto

@@ -51,7 +51,13 @@ message Package {
 }
 
 message PerClassMetadata {
-    repeated string annotation_class_names = 1;
+  repeated string annotation_class_names = 1;
+  // Not all methods will be present here, only those with something interesting to report.
+  map<string, PerMethodMetadata> per_method_metadata = 2;
+}
+
+message PerMethodMetadata {
+  repeated string annotation_class_names = 1;
 }
 
 service Lifecycle {

java/gazelle/testdata/attribute_setting/src/test/com/example/myproject/BUILD.out

@@ -15,6 +15,22 @@ java_test_suite(
     ],
 )
 
+java_junit5_test(
+    name = "MixedTest",
+    srcs = ["MixedTest.java"],
+    flaky = True,
+    test_class = "com.example.myproject.MixedTest",
+    runtime_deps = [
+        "@maven//:org_junit_jupiter_junit_jupiter_engine",
+        "@maven//:org_junit_platform_junit_platform_launcher",
+        "@maven//:org_junit_platform_junit_platform_reporting",
+    ],
+    deps = [
+        "//src/main/com/example/annotation",
+        "@maven//:org_junit_jupiter_junit_jupiter_api",
+    ],
+)
+
 java_junit5_test(
     name = "RandomTest",
     srcs = ["RandomTest.java"],

java/gazelle/testdata/attribute_setting/src/test/com/example/myproject/MixedTest.java

@@ -0,0 +1,25 @@
+package com.example.myproject;
+
+import com.example.annotation.FlakyTest;
+import org.junit.jupiter.api.Test;
+
+import java.util.Random;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class MixedTest {
+    @Test
+    @FlakyTest
+    public void unreliableTest() {
+        Random random = new Random();
+        int r = random.nextInt(2);
+        assertEquals(r, 0);
+    }
+
+    @Test
+    public void reliableTest() {
+        Random random = new Random();
+        int r = random.nextInt(2);
+        assertTrue(r < 2);
+    }
+}

java/src/com/github/bazel_contrib/contrib_rules_jvm/javaparser/generators/ClasspathParser.java

@@ -5,9 +5,7 @@
 import static javax.lang.model.element.Modifier.STATIC;
 
 import com.google.common.base.Splitter;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Lists;
 import com.sun.source.tree.AnnotationTree;
 import com.sun.source.tree.ArrayTypeTree;
@@ -31,7 +29,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
+import java.util.SortedMap;
 import java.util.SortedSet;
 import java.util.TreeMap;
 import java.util.TreeSet;
@@ -57,7 +57,7 @@ public class ClasspathParser {
 
   // Mapping from fully-qualified class-name to class-names of annotations on that class.
   // Annotations will be fully-qualified where that's known, and not where not known.
-  private final Map<String, SortedSet<String>> annotatedClasses = new TreeMap<>();
+  final Map<String, PerClassData> perClassData = new TreeMap<>();
 
   // get the system java compiler instance
   private static final JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
@@ -67,6 +67,46 @@ public ClasspathParser() {
     // Doesn't need to do anything currently
   }
 
+  static class PerClassData {
+    public PerClassData() {
+      this(new TreeSet<>(), new TreeMap<>());
+    }
+
+    @Override
+    public String toString() {
+      return "PerClassData{"
+          + "annotations="
+          + annotations
+          + ", perMethodAnnotations="
+          + perMethodAnnotations
+          + '}';
+    }
+
+    public PerClassData(
+        SortedSet<String> annotations, SortedMap<String, SortedSet<String>> perMethodAnnotations) {
+      this.annotations = annotations;
+      this.perMethodAnnotations = perMethodAnnotations;
+    }
+
+    final SortedSet<String> annotations;
+
+    final SortedMap<String, SortedSet<String>> perMethodAnnotations;
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      PerClassData that = (PerClassData) o;
+      return Objects.equals(annotations, that.annotations)
+          && Objects.equals(perMethodAnnotations, that.perMethodAnnotations);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(annotations, perMethodAnnotations);
+    }
+  }
+
   public ImmutableSet<String> getUsedTypes() {
     return ImmutableSet.copyOf(usedTypes);
   }
@@ -87,15 +127,6 @@ public ImmutableSet<String> getMainClasses() {
     return ImmutableSet.copyOf(mainClasses);
   }
 
-  public ImmutableMap<String, ImmutableSortedSet<String>> getAnnotatedClasses() {
-    ImmutableMap.Builder<String, ImmutableSortedSet<String>> builder =
-        ImmutableMap.builderWithExpectedSize(annotatedClasses.size());
-    for (Map.Entry<String, SortedSet<String>> entry : annotatedClasses.entrySet()) {
-      builder.put(entry.getKey(), ImmutableSortedSet.copyOf(entry.getValue()));
-    }
-    return builder.build();
-  }
-
   public void parseClasses(Path directory, List<String> files) throws IOException {
     StandardJavaFileManager fileManager = compiler.getStandardFileManager(null, null, null);
     List<? extends JavaFileObject> objectFiles =
@@ -245,6 +276,20 @@ public Void visitMethod(com.sun.source.tree.MethodTree m, Void v) {
         checkFullyQualifiedType(param.getType());
         handleAnnotations(param.getModifiers().getAnnotations());
       }
+
+      for (AnnotationTree annotation : m.getModifiers().getAnnotations()) {
+        String annotationClassName = annotation.getAnnotationType().toString();
+        String importedFullyQualified = currentFileImports.get(annotationClassName);
+        String currentFullyQualifiedClass = fullyQualify(currentPackage, currentClassName);
+        if (importedFullyQualified != null) {
+          noteAnnotatedMethod(
+              currentFullyQualifiedClass, m.getName().toString(), importedFullyQualified);
+        } else {
+          noteAnnotatedMethod(
+              currentFullyQualifiedClass, m.getName().toString(), annotationClassName);
+        }
+      }
+
       return super.visitMethod(m, v);
     }
 
@@ -333,10 +378,27 @@ private Set<String> checkFullyQualifiedType(Tree identifier) {
 
   private void noteAnnotatedClass(
       String annotatedFullyQualifiedClassName, String annotationFullyQualifiedClassName) {
-    if (!annotatedClasses.containsKey(annotatedFullyQualifiedClassName)) {
-      annotatedClasses.put(annotatedFullyQualifiedClassName, new TreeSet<>());
+    if (!perClassData.containsKey(annotatedFullyQualifiedClassName)) {
+      perClassData.put(annotatedFullyQualifiedClassName, new PerClassData());
+    }
+    perClassData
+        .get(annotatedFullyQualifiedClassName)
+        .annotations
+        .add(annotationFullyQualifiedClassName);
+  }
+
+  private void noteAnnotatedMethod(
+      String annotatedFullyQualifiedClassName,
+      String methodName,
+      String annotationFullyQualifiedClassName) {
+    if (!perClassData.containsKey(annotatedFullyQualifiedClassName)) {
+      perClassData.put(annotatedFullyQualifiedClassName, new PerClassData());
+    }
+    PerClassData data = perClassData.get(annotatedFullyQualifiedClassName);
+    if (!data.perMethodAnnotations.containsKey(methodName)) {
+      data.perMethodAnnotations.put(methodName, new TreeSet<>());
     }
-    annotatedClasses.get(annotatedFullyQualifiedClassName).add(annotationFullyQualifiedClassName);
+    data.perMethodAnnotations.get(methodName).add(annotationFullyQualifiedClassName);
   }
 
   private String fullyQualify(String packageName, String className) {

java/src/com/github/bazel_contrib/contrib_rules_jvm/javaparser/generators/GrpcServer.java

@@ -7,9 +7,9 @@
 import com.gazelle.java.javaparser.v0.Package.Builder;
 import com.gazelle.java.javaparser.v0.ParsePackageRequest;
 import com.gazelle.java.javaparser.v0.PerClassMetadata;
+import com.gazelle.java.javaparser.v0.PerMethodMetadata;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Iterables;
 import io.grpc.Server;
 import io.grpc.ServerBuilder;
@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.SortedSet;
 import java.util.concurrent.TimeUnit;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -154,13 +155,20 @@ private Package getImports(ParsePackageRequest request) {
               .addAllImportedPackagesWithoutSpecificClasses(
                   parser.getUsedPackagesWithoutSpecificTypes())
               .addAllMains(parser.getMainClasses());
-      for (Map.Entry<String, ImmutableSortedSet<String>> annotations :
-          parser.getAnnotatedClasses().entrySet()) {
-        packageBuilder.putPerClassMetadata(
-            annotations.getKey(),
+      for (Map.Entry<String, ClasspathParser.PerClassData> classEntry :
+          parser.perClassData.entrySet()) {
+        PerClassMetadata.Builder perClassMetadata =
             PerClassMetadata.newBuilder()
-                .addAllAnnotationClassNames(annotations.getValue())
-                .build());
+                .addAllAnnotationClassNames(classEntry.getValue().annotations);
+        for (Map.Entry<String, SortedSet<String>> methodEntry :
+            classEntry.getValue().perMethodAnnotations.entrySet()) {
+          perClassMetadata.putPerMethodMetadata(
+              methodEntry.getKey(),
+              PerMethodMetadata.newBuilder()
+                  .addAllAnnotationClassNames(methodEntry.getValue())
+                  .build());
+        }
+        packageBuilder.putPerClassMetadata(classEntry.getKey(), perClassMetadata.build());
       }
       return packageBuilder.build();
     }

java/test/com/github/bazel_contrib/contrib_rules_jvm/javaparser/generators/ClasspathParserTest.java

@@ -14,6 +14,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.SortedSet;
+import java.util.TreeMap;
 import java.util.TreeSet;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -191,8 +193,26 @@ public void testAnnotationAfterImport() throws IOException {
     assertEquals(
         Map.of(
             "workspace.com.gazelle.java.javaparser.generators.AnnotationAfterImport",
-            treeSet("com.example.FlakyTest")),
-        parser.getAnnotatedClasses());
+            new ClasspathParser.PerClassData(treeSet("com.example.FlakyTest"), new TreeMap<>())),
+        parser.perClassData);
+  }
+
+  @Test
+  public void testAnnotationAfterImportOnMethod() throws IOException {
+    List<? extends JavaFileObject> files =
+        List.of(
+            testFiles.get(
+                "/workspace/com/gazelle/java/javaparser/generators/AnnotationAfterImportOnMethod.java"));
+    parser.parseClasses(files);
+
+    TreeMap<String, SortedSet<String>> expectedPerMethodAnnotations = new TreeMap<>();
+    expectedPerMethodAnnotations.put("someTest", treeSet("org.junit.jupiter.api.Test"));
+
+    assertEquals(
+        Map.of(
+            "workspace.com.gazelle.java.javaparser.generators.AnnotationAfterImportOnMethod",
+            new ClasspathParser.PerClassData(new TreeSet<>(), expectedPerMethodAnnotations)),
+        parser.perClassData);
   }
 
   @Test
@@ -208,8 +228,8 @@ public void testAnnotationFromJavaStandardLibrary() throws IOException {
     assertEquals(
         Map.of(
             "workspace.com.gazelle.java.javaparser.generators.AnnotationFromJavaStandardLibrary",
-            treeSet("Deprecated")),
-        parser.getAnnotatedClasses());
+            new ClasspathParser.PerClassData(treeSet("Deprecated"), new TreeMap<>())),
+        parser.perClassData);
   }
 
   @Test
@@ -225,8 +245,8 @@ public void testAnnotationWithoutImport() throws IOException {
     assertEquals(
         Map.of(
             "workspace.com.gazelle.java.javaparser.generators.AnnotationWithoutImport",
-            treeSet("WhoKnowsWhereIAmFrom")),
-        parser.getAnnotatedClasses());
+            new ClasspathParser.PerClassData(treeSet("WhoKnowsWhereIAmFrom"), new TreeMap<>())),
+        parser.perClassData);
   }
 
   @Test