Support entitlements in internal cluster tests

Author: mosche

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/ElasticsearchTestBasePlugin.java

@@ -34,6 +34,7 @@
 import java.io.File;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Stream;
 
 import javax.inject.Inject;
@@ -50,6 +51,8 @@ public abstract class ElasticsearchTestBasePlugin implements Plugin<Project> {
 
     public static final String DUMP_OUTPUT_ON_FAILURE_PROP_NAME = "dumpOutputOnFailure";
 
+    public static final Set<String> TEST_TASKS_WITH_ENTITLEMENTS = Set.of("test", "internalClusterTest");
+
     @Inject
     protected abstract ProviderFactory getProviderFactory();
 
@@ -178,10 +181,17 @@ public void execute(Task t) {
             SourceSetContainer sourceSets = project.getExtensions().getByType(SourceSetContainer.class);
             SourceSet mainSourceSet = sourceSets.findByName(SourceSet.MAIN_SOURCE_SET_NAME);
             SourceSet testSourceSet = sourceSets.findByName(SourceSet.TEST_SOURCE_SET_NAME);
-            if ("test".equals(test.getName()) && mainSourceSet != null && testSourceSet != null) {
+            SourceSet internalClusterTestSourceSet = sourceSets.findByName("internalClusterTest");
+
+            if (TEST_TASKS_WITH_ENTITLEMENTS.contains(test.getName()) && mainSourceSet != null && testSourceSet != null) {
                 FileCollection mainRuntime = mainSourceSet.getRuntimeClasspath();
                 FileCollection testRuntime = testSourceSet.getRuntimeClasspath();
-                FileCollection testOnlyFiles = testRuntime.minus(mainRuntime);
+                FileCollection internalClusterTestRuntime = "internalClusterTest".equals(test.getName())
+                    && internalClusterTestSourceSet == null
+                        ? project.files() // empty file collection
+                        : internalClusterTestSourceSet.getRuntimeClasspath();
+                FileCollection testOnlyFiles = testRuntime.plus(internalClusterTestRuntime).minus(mainRuntime);
+
                 test.doFirst(task -> test.environment("es.entitlement.testOnlyPath", testOnlyFiles.getAsPath()));
             }
 
@@ -241,16 +251,18 @@ public void execute(Task t) {
      * Computes and sets the {@code --patch-module=java.base} and {@code --add-opens=java.base} JVM command line options.
      */
     private void configureJavaBaseModuleOptions(Project project) {
-        project.getTasks().withType(Test.class).matching(task -> task.getName().equals("test")).configureEach(test -> {
-            FileCollection patchedImmutableCollections = patchedImmutableCollections(project);
+        project.getTasks().withType(Test.class).configureEach(test -> {
+            // patch immutable collections only for "test" task
+            FileCollection patchedImmutableCollections = test.getName().equals("test") ? patchedImmutableCollections(project) : null;
             if (patchedImmutableCollections != null) {
                 test.getInputs().files(patchedImmutableCollections);
                 test.systemProperty("tests.hackImmutableCollections", "true");
             }
 
-            FileCollection entitlementBridge = entitlementBridge(project);
+            FileCollection entitlementBridge = TEST_TASKS_WITH_ENTITLEMENTS.contains(test.getName()) ? entitlementBridge(project) : null;
             if (entitlementBridge != null) {
                 test.getInputs().files(entitlementBridge);
+                test.systemProperty("es.entitlement.enableForTests", "true");
             }
 
             test.getJvmArgumentProviders().add(() -> {
@@ -312,7 +324,9 @@ private static void configureEntitlements(Project project) {
         }
         FileCollection bridgeFiles = bridgeConfig;
 
-        project.getTasks().withType(Test.class).configureEach(test -> {
+        project.getTasks().withType(Test.class)
+            .matching(test -> TEST_TASKS_WITH_ENTITLEMENTS.contains(test.getName()))
+            .configureEach(test -> {
             // See also SystemJvmOptions.maybeAttachEntitlementAgent.
 
             // Agent

build-tools/src/main/java/org/elasticsearch/gradle/test/TestBuildInfoPlugin.java

@@ -56,11 +56,5 @@ public void apply(Project project) {
         project.getTasks().withType(ProcessResources.class).named("processResources").configure(task -> {
             task.into("META-INF", copy -> copy.from(testBuildInfoTask));
         });
-
-        if (project.getRootProject().getName().equals("elasticsearch")) {
-            project.getTasks().withType(Test.class).matching(test -> List.of("test").contains(test.getName())).configureEach(test -> {
-                test.systemProperty("es.entitlement.enableForTests", "true");
-            });
-        }
     }
 }

libs/build.gradle

@@ -46,13 +46,13 @@ configure(childProjects.values()) {
     apply plugin: 'elasticsearch.build'
   }
 
-  // This is for any code potentially included in the server at runtime.
-  // Omit oddball libraries that aren't in server.
-  def nonServerLibs = ['plugin-scanner']
-  if (false == nonServerLibs.contains(project.name)) {
-    project.getTasks().withType(Test.class).matching(test -> ['test'].contains(test.name)).configureEach(test -> {
-      test.systemProperty('es.entitlement.enableForTests', 'true')
-    })
-  }
+//  // This is for any code potentially included in the server at runtime.
+//  // Omit oddball libraries that aren't in server.
+//  def nonServerLibs = ['plugin-scanner']
+//  if (false == nonServerLibs.contains(project.name)) {
+//    project.getTasks().withType(Test.class).matching(test -> ['test', 'internalClusterTest'].contains(test.name)).configureEach(test -> {
+//      test.systemProperty('es.entitlement.enableForTests', 'true')
+//    })
+//  }
 
 }