Look for annotations on methods as well as classes

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/private/sorted_multiset/multiset.go

@@ -33,7 +33,15 @@ func (s *SortedMultiSet[K, V]) Keys() []K {
 	return s.keys.SortedSlice()
 }
 
-func (s *SortedMultiSet[K, V]) Values(key K) []V {
+func (s *SortedMultiSet[K, V]) Values(key K) *sorted_set.SortedSet[V] {
+	if s == nil {
+		return sorted_set.NewSortedSet[V](nil)
+	}
+
+	return s.ms[key]
+}
+
+func (s *SortedMultiSet[K, V]) SortedValues(key K) []V {
 	if s == nil {
 		return nil
 	}

java/gazelle/private/sorted_multiset/multiset_test.go

@@ -15,7 +15,7 @@ func TestNil(t *testing.T) {
 		t.Errorf("want len(nil.Keys()) == 0, got %d", gotKeysLen)
 	}
 
-	gotValues := s.Values("foo")
+	gotValues := s.SortedValues("foo")
 	if gotValuesLen := len(gotValues); gotValuesLen != 0 {
 		t.Errorf("want len(nil.Values(\"foo\")) == 0, got %d", gotValuesLen)
 	}
@@ -30,7 +30,12 @@ func TestMultiSet(t *testing.T) {
 			t.Errorf("want no keys, got %d: %v", gotKeysLen, gotKeys)
 		}
 
-		gotValues := s.Values("tasty")
+		gotValuesSlice := s.SortedValues("tasty")
+		if gotValuesLen := len(gotValuesSlice); gotValuesLen != 0 {
+			t.Errorf("want no values for tasty, got %d: %v", gotValuesLen, gotValuesSlice)
+		}
+
+		gotValues := s.Values("tasty").SortedSlice()
 		if gotValuesLen := len(gotValues); gotValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotValuesLen, gotValues)
 		}
@@ -45,12 +50,12 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if gotBadValuesLen := len(gotBadValues); gotBadValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotBadValuesLen, gotBadValues)
 		}
@@ -65,12 +70,12 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"cheese", "hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if gotBadValuesLen := len(gotBadValues); gotBadValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotBadValuesLen, gotBadValues)
 		}
@@ -85,12 +90,12 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"cheese", "hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if gotBadValuesLen := len(gotBadValues); gotBadValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotBadValuesLen, gotBadValues)
 		}
@@ -105,13 +110,18 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"cheese", "hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
+		gotTastySortedSetValues := s.Values("tasty")
+		if !reflect.DeepEqual(wantTastyValues, gotTastySortedSetValues.SortedSlice()) {
+			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
+		}
+
 		wantBadValues := []string{"soil"}
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if !reflect.DeepEqual(wantBadValues, gotBadValues) {
 			t.Errorf("want bad values %v got %v", wantBadValues, gotBadValues)
 		}

java/gazelle/private/sorted_set/btreeset.go

@@ -66,6 +66,9 @@ func (s *SortedSet[T]) Add(v T) {
 
 // AddAll adds all of the elements from the other SortedSet to this one.
 func (s *SortedSet[T]) AddAll(other *SortedSet[T]) {
+	if other == nil {
+		return
+	}
 	other.tree.Ascend(func(v T) bool {
 		s.Add(v)
 		return true

java/gazelle/private/sorted_set/btreeset_test.go

@@ -251,6 +251,13 @@ func TestAddAll(t *testing.T) {
 	if !reflect.DeepEqual(want, got) {
 		t.Errorf("want %v got %v", want, got)
 	}
+
+	s.AddAll(nil)
+
+	got = s.SortedSlice()
+	if !reflect.DeepEqual(want, got) {
+		t.Errorf("want %v got %v", want, got)
+	}
 }
 
 func TestFilter(t *testing.T) {

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);
+    }
+}