Allow setting wrapper targets based on annotations

This allows for registering that, e.g.

```java
@RequiresNetwork
class SomeTest {}
```

should be generated as:

```starlark
requires_network(
    java_test,
    ...
)
```

instead of:

```starlark
java_test(
    ...
)
```

In suite mode, separate targets will be generated for these special
targets.

java/gazelle/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "//java/gazelle/private/javaparser",
         "//java/gazelle/private/logconfig",
         "//java/gazelle/private/maven",
+        "//java/gazelle/private/sorted_multiset",
         "//java/gazelle/private/sorted_set",
         "@bazel_gazelle//config:go_default_library",
         "@bazel_gazelle//label:go_default_library",

java/gazelle/configure.go

@@ -20,17 +20,20 @@ import (
 type Configurer struct {
 	lang                  *javaLang
 	annotationToAttribute annotationToAttribute
+	annotationToWrapper   annotationToWrapper
 }
 
 func NewConfigurer(lang *javaLang) *Configurer {
 	return &Configurer{
 		lang:                  lang,
 		annotationToAttribute: make(annotationToAttribute),
+		annotationToWrapper:   make(annotationToWrapper),
 	}
 }
 
 func (jc Configurer) RegisterFlags(fs *flag.FlagSet, cmd string, c *config.Config) {
 	fs.Var(&jc.annotationToAttribute, "java-annotation-to-attribute", "Mapping of annotations (on test classes) to attributes which should be set for that test rule. Examples: com.example.annotations.FlakyTest=flaky=True com.example.annotations.SlowTest=timeout=\"long\"")
+	fs.Var(&jc.annotationToWrapper, "java-annotation-to-wrapper", "Mapping of annotations (on test classes) to wrapper rules which should be used around the test rule. Example: com.example.annotations.RequiresNetwork=@some//wrapper:file.bzl=requires_network")
 }
 
 func (jc *Configurer) CheckFlags(fs *flag.FlagSet, c *config.Config) error {
@@ -40,6 +43,9 @@ func (jc *Configurer) CheckFlags(fs *flag.FlagSet, c *config.Config) error {
 			cfgs[""].MapAnnotationToAttribute(annotation, k, v)
 		}
 	}
+	for annotation, wrapper := range jc.annotationToWrapper {
+		cfgs[""].MapAnnotationToWrapper(annotation, wrapper.symbol)
+	}
 	return nil
 }
 
@@ -159,3 +165,47 @@ func (f *annotationToAttribute) Set(value string) error {
 	(*f)[annotationClassName][key] = parsedValue
 	return nil
 }
+
+type loadInfo struct {
+	from   string
+	symbol string
+}
+
+type annotationToWrapper map[string]loadInfo
+
+func (f *annotationToWrapper) String() string {
+	s := "annotationToWrapper{"
+	for a, li := range *f {
+		s += a + ": "
+		s += fmt.Sprintf(`load("%s", "%s")`, li.from, li.symbol)
+	}
+	s += "}"
+	return s
+}
+
+func (f *annotationToWrapper) Set(value string) error {
+	parts := strings.Split(value, "=")
+	if len(parts) != 2 {
+		return fmt.Errorf("want --java-annotation-to-wrapper to have format com.example.RequiresNetwork=@some_repo//has:wrapper.bzl,wrapper_rule but didn't see exactly one equals sign")
+	}
+	annotation := parts[0]
+
+	if _, ok := (*f)[annotation]; ok {
+		return fmt.Errorf("saw conflicting values for --java-annotation-to-wrapper flag for annotation %v", annotation)
+	}
+
+	vParts := strings.Split(parts[1], ",")
+	if len(vParts) != 2 {
+		return fmt.Errorf("want --java-annotation-to-wrapper to have format com.example.RequiresNetwork=@some_repo//has:wrapper.bzl,wrapper_rule but didn't see exactly one comma after equals sign")
+	}
+
+	from := vParts[0]
+	symbol := vParts[1]
+
+	(*f)[annotation] = loadInfo{
+		from:   from,
+		symbol: symbol,
+	}
+
+	return nil
+}

java/gazelle/generate.go

@@ -49,6 +49,11 @@ func javaFileLess(l, r javaFile) bool {
 	return l.pathRelativeToBazelWorkspaceRoot < r.pathRelativeToBazelWorkspaceRoot
 }
 
+type separateJavaTestReasons struct {
+	attributes map[string]bzl.Expr
+	wrapper    string
+}
+
 // GenerateRules extracts build metadata from source files in a directory.
 //
 // See language.GenerateRules for more information.
@@ -171,7 +176,7 @@ func (l javaLang) GenerateRules(args language.GenerateArgs) language.GenerateRes
 	testJavaImports := sorted_set.NewSortedSet([]string{})
 
 	// Java Test files which need to be generated separately from any others because they have explicit attribute overrides.
-	separateTestJavaFiles := make(map[javaFile]map[string]bzl.Expr)
+	separateTestJavaFiles := make(map[javaFile]separateJavaTestReasons)
 
 	// Files which are used by non-test classes in test java packages.
 	testHelperJavaFiles := sorted_set.NewSortedSetFn([]javaFile{}, javaFileLess)
@@ -289,8 +294,8 @@ func (l javaLang) GenerateRules(args language.GenerateArgs) language.GenerateRes
 		switch cfg.TestMode() {
 		case "file":
 			for _, tf := range testJavaFiles.SortedSlice() {
-				extraAttributes := separateTestJavaFiles[tf]
-				generateJavaTest(args.Rel, tf, isModule, testJavaImportsWithHelpers, extraAttributes, &res)
+				separateJavaTestReasons := separateTestJavaFiles[tf]
+				generateJavaTest(args.Rel, tf, isModule, testJavaImportsWithHelpers, separateJavaTestReasons.wrapper, separateJavaTestReasons.attributes, &res)
 			}
 
 		case "suite":
@@ -323,7 +328,8 @@ func (l javaLang) GenerateRules(args language.GenerateArgs) language.GenerateRes
 				sortedSeparateTestJavaFiles.Add(src)
 			}
 			for _, src := range sortedSeparateTestJavaFiles.SortedSlice() {
-				generateJavaTest(args.Rel, src, isModule, testJavaImportsWithHelpers, separateTestJavaFiles[src], &res)
+				separateJavaTestReasons := separateTestJavaFiles[src]
+				generateJavaTest(args.Rel, src, isModule, testJavaImportsWithHelpers, separateJavaTestReasons.wrapper, separateJavaTestReasons.attributes, &res)
 			}
 		}
 	}
@@ -356,10 +362,11 @@ func addNonLocalImports(to *sorted_set.SortedSet[string], from *sorted_set.Sorte
 	}
 }
 
-func accumulateJavaFile(cfg *javaconfig.Config, testJavaFiles, testHelperJavaFiles *sorted_set.SortedSet[javaFile], separateTestJavaFiles map[javaFile]map[string]bzl.Expr, file javaFile, perClassMetadata map[string]java.PerClassMetadata, log zerolog.Logger) {
+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 maven.IsTestFile(filepath.Base(file.pathRelativeToBazelWorkspaceRoot)) {
 		annotationClassNames := perClassMetadata[file.FullyQualifiedClassName()].AnnotationClassNames
 		perFileAttrs := make(map[string]bzl.Expr)
+		wrapper := ""
 		for _, annotationClassName := range annotationClassNames.SortedSlice() {
 			if attrs, ok := cfg.AttributesForAnnotation(annotationClassName); ok {
 				for k, v := range attrs {
@@ -369,11 +376,21 @@ func accumulateJavaFile(cfg *javaconfig.Config, testJavaFiles, testHelperJavaFil
 					perFileAttrs[k] = v
 				}
 			}
+			newWrapper, ok := cfg.WrapperForAnnotation(annotationClassName)
+			if ok {
+				if wrapper != "" {
+					log.Error().Str("file", file.pathRelativeToBazelWorkspaceRoot).Msgf("Saw conflicting wrappers from annotations: %v and %v. Picking one at random.", wrapper, newWrapper)
+				}
+				wrapper = newWrapper
+			}
 		}
-		if len(perFileAttrs) == 0 {
+		if len(perFileAttrs) == 0 && wrapper == "" {
 			testJavaFiles.Add(file)
 		} else {
-			separateTestJavaFiles[file] = perFileAttrs
+			separateTestJavaFiles[file] = separateJavaTestReasons{
+				attributes: perFileAttrs,
+				wrapper:    wrapper,
+			}
 		}
 	} else {
 		testHelperJavaFiles.Add(file)
@@ -409,7 +426,7 @@ func generateJavaBinary(m main, libName string, res *language.GenerateResult) {
 	res.Imports = append(res.Imports, []string{})
 }
 
-func generateJavaTest(pathToPackageRelativeToBazelWorkspace string, f javaFile, includePackageInName bool, imports *sorted_set.SortedSet[string], extraAttributes map[string]bzl.Expr, res *language.GenerateResult) {
+func generateJavaTest(pathToPackageRelativeToBazelWorkspace string, f javaFile, includePackageInName bool, imports *sorted_set.SortedSet[string], wrapper string, extraAttributes map[string]bzl.Expr, res *language.GenerateResult) {
 	fullyQualifiedTestClass := f.FullyQualifiedClassName()
 	var testName string
 	if includePackageInName {
@@ -418,14 +435,23 @@ func generateJavaTest(pathToPackageRelativeToBazelWorkspace string, f javaFile,
 		testName = f.ClassName()
 	}
 
-	ruleKind := "java_test"
+	javaRuleKind := "java_test"
 	var runtimeDeps []string
 	if importsJunit5(imports) {
-		ruleKind = "java_junit5_test"
+		javaRuleKind = "java_junit5_test"
 		runtimeDeps = append(runtimeDeps, mapStringSlice(junit5RuntimeDeps, maven.LabelFromArtifact)...)
 	}
 
+	ruleKind := javaRuleKind
+	if wrapper != "" {
+		ruleKind = wrapper
+	}
+
 	r := rule.NewRule(ruleKind, testName)
+	if wrapper != "" {
+		r.AddArg(&bzl.Ident{Name: javaRuleKind})
+	}
+
 	path := strings.TrimPrefix(f.pathRelativeToBazelWorkspaceRoot, pathToPackageRelativeToBazelWorkspace+"/")
 	r.SetAttr("srcs", []string{path})
 	r.SetAttr("test_class", fullyQualifiedTestClass)

java/gazelle/generate_test.go

@@ -110,7 +110,7 @@ func TestSingleJavaTestFile(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			var res language.GenerateResult
 
-			generateJavaTest("", f, tc.includePackageInName, sorted_set.NewSortedSet(tc.imports), nil, &res)
+			generateJavaTest("", f, tc.includePackageInName, sorted_set.NewSortedSet(tc.imports), "", nil, &res)
 
 			require.Len(t, res.Gen, 1, "want 1 generated rule")
 

java/gazelle/javaconfig/config.go

@@ -44,6 +44,7 @@ func (c *Config) NewChild() *Config {
 		repoRoot:              c.repoRoot,
 		testMode:              c.testMode,
 		annotationToAttribute: c.annotationToAttribute,
+		annotationToWrapper:   c.annotationToWrapper,
 	}
 }
 
@@ -68,6 +69,7 @@ type Config struct {
 	repoRoot              string
 	testMode              string
 	annotationToAttribute map[string]map[string]bzl.Expr
+	annotationToWrapper   map[string]string
 }
 
 type LoadInfo struct {
@@ -85,6 +87,7 @@ func New(repoRoot string) *Config {
 		repoRoot:              repoRoot,
 		testMode:              "suite",
 		annotationToAttribute: make(map[string]map[string]bzl.Expr),
+		annotationToWrapper:   make(map[string]string),
 	}
 }
 
@@ -154,3 +157,12 @@ func (c *Config) AttributesForAnnotation(annotation string) (map[string]bzl.Expr
 	m, ok := c.annotationToAttribute[annotation]
 	return m, ok
 }
+
+func (c *Config) MapAnnotationToWrapper(annotation string, wrapper string) {
+	c.annotationToWrapper[annotation] = wrapper
+}
+
+func (c *Config) WrapperForAnnotation(annotation string) (string, bool) {
+	s, ok := c.annotationToWrapper[annotation]
+	return s, ok
+}

java/gazelle/lang.go

@@ -7,6 +7,7 @@ import (
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/javaparser"
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/logconfig"
 	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/maven"
+	"github.com/bazel-contrib/rules_jvm/java/gazelle/private/sorted_multiset"
 	"github.com/bazelbuild/bazel-gazelle/config"
 	"github.com/bazelbuild/bazel-gazelle/language"
 	"github.com/bazelbuild/bazel-gazelle/resolve"
@@ -75,7 +76,7 @@ var kindWithoutRuntimeDeps = rule.KindInfo{
 }
 
 func (l javaLang) Kinds() map[string]rule.KindInfo {
-	return map[string]rule.KindInfo{
+	kinds := map[string]rule.KindInfo{
 		"java_binary":        kindWithRuntimeDeps,
 		"java_junit5_test":   kindWithRuntimeDeps,
 		"java_library":       kindWithRuntimeDeps,
@@ -84,9 +85,16 @@ func (l javaLang) Kinds() map[string]rule.KindInfo {
 		"java_proto_library": kindWithoutRuntimeDeps,
 		"java_grpc_library":  kindWithoutRuntimeDeps,
 	}
+
+	c := l.Configurer.(*Configurer)
+	for _, wrapper := range c.annotationToWrapper {
+		kinds[wrapper.symbol] = kindWithRuntimeDeps
+	}
+
+	return kinds
 }
 
-var javaLoads = []rule.LoadInfo{
+var baseJavaLoads = []rule.LoadInfo{
 	{
 		Name: "@io_grpc_grpc_java//:java_grpc_library.bzl",
 		Symbols: []string{
@@ -112,7 +120,30 @@ var javaLoads = []rule.LoadInfo{
 }
 
 func (l javaLang) Loads() []rule.LoadInfo {
-	return javaLoads
+	c := l.Configurer.(*Configurer)
+	if len(c.annotationToWrapper) == 0 {
+		return baseJavaLoads
+	}
+
+	s := sorted_multiset.NewSortedMultiSet[string, string]()
+	for _, li := range baseJavaLoads {
+		for _, symbol := range li.Symbols {
+			s.Add(li.Name, symbol)
+		}
+	}
+
+	for _, wrapper := range c.annotationToWrapper {
+		s.Add(wrapper.from, wrapper.symbol)
+	}
+
+	var loads []rule.LoadInfo
+	for _, name := range s.Keys() {
+		loads = append(loads, rule.LoadInfo{
+			Name:    name,
+			Symbols: s.Values(name),
+		})
+	}
+	return loads
 }
 
 func (l javaLang) Fix(c *config.Config, f *rule.File) {}

java/gazelle/testdata/annotation_to_wrapper/arguments.txt

@@ -0,0 +1 @@
+-java-annotation-to-wrapper=com.example.annotations.RequiresNetwork=@some//:requires_network.bzl,requires_network

java/gazelle/testdata/annotation_to_wrapper/expectedStderr.txt

@@ -0,0 +1,3 @@
+12:00AM WRN java/gazelle/private/maven/resolver.go:33 > not loading maven dependencies error="open %WORKSPACEPATH%/maven_install.json: no such file or directory" _c=maven-resolver
+12:00AM WRN java/gazelle/resolve.go:215 > missing import com.example.tests pkg=com.example.tests targets=[]
+12:00AM WRN java/gazelle/resolve.go:215 > missing import org.junit.Test pkg=org.junit targets=[]

java/gazelle/testdata/annotation_to_wrapper/src/main/com/example/annotations/BUILD.out

@@ -0,0 +1,7 @@
+load("@rules_java//java:defs.bzl", "java_library")
+
+java_library(
+    name = "annotations",
+    srcs = ["RequiresNetwork.java"],
+    visibility = ["//:__subpackages__"],
+)

java/gazelle/testdata/annotation_to_wrapper/src/main/com/example/annotations/RequiresNetwork.java

@@ -0,0 +1,3 @@
+package com.example.annotations;
+
+public @interface RequiresNetwork {}

java/gazelle/testdata/annotation_to_wrapper/src/test/java/com/example/tests/BUILD.out

@@ -0,0 +1,17 @@
+load("@contrib_rules_jvm//java:defs.bzl", "java_test_suite")
+load("@rules_java//java:defs.bzl", "java_test")
+load("@some//:requires_network.bzl", "requires_network")
+
+java_test_suite(
+    name = "tests",
+    srcs = ["SimpleTest.java"],
+    deps = ["//src/main/com/example/annotations"],
+)
+
+requires_network(
+    java_test,
+    name = "RequiresNetworkTest",
+    srcs = ["RequiresNetworkTest.java"],
+    test_class = "com.example.tests.RequiresNetworkTest",
+    deps = ["//src/main/com/example/annotations"],
+)

java/gazelle/testdata/annotation_to_wrapper/src/test/java/com/example/tests/RequiresNetworkTest.java

@@ -0,0 +1,12 @@
+package com.example.tests;
+
+import com.example.annotations.RequiresNetwork;
+
+import org.junit.Test;
+
+@RequiresNetwork
+public class RequiresNetworkTest {
+  @Test
+  public void passes() {}
+}
+

java/gazelle/testdata/annotation_to_wrapper/src/test/java/com/example/tests/SimpleTest.java

@@ -0,0 +1,8 @@
+package com.example.tests;
+
+import org.junit.Test;
+
+public class SimpleTest {
+  @Test
+  public void passes() {}
+}

repositories.bzl

@@ -59,10 +59,10 @@ def contrib_rules_jvm_gazelle_deps():
     maybe(
         http_archive,
         name = "bazel_gazelle",
-        sha256 = "b794b5c4af78574e4bf574c6b2336bbffbdfeadfaae1c14c6184f154e44fdb47",
-        strip_prefix = "bazel-gazelle-622d8889c63227a71d6b393ba5e3a8b8a6761466",
+        sha256 = "68f2caeb8849ac0bcd257c801c435c488ad282695e7888aeffc3dfc707803ba6",
+        strip_prefix = "bazel-gazelle-cfbfdff7ff0ed8913d9a24648709ae1aca784436",
         urls = [
-            "https://github.com/bazelbuild/bazel-gazelle/archive/622d8889c63227a71d6b393ba5e3a8b8a6761466.tar.gz",
+            "https://github.com/bazelbuild/bazel-gazelle/archive/cfbfdff7ff0ed8913d9a24648709ae1aca784436.tar.gz",
         ],
         patch_args = ["-p1"],
         patches = [