Allow overriding of plugin metadata files in integration tests (#120245) (#120297)

Author: mark-vieira

test/test-clusters/src/main/java/org/elasticsearch/test/cluster/local/AbstractLocalClusterFactory.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.test.cluster.local.distribution.DistributionResolver;
 import org.elasticsearch.test.cluster.local.distribution.DistributionType;
 import org.elasticsearch.test.cluster.local.model.User;
+import org.elasticsearch.test.cluster.util.ArchivePatcher;
 import org.elasticsearch.test.cluster.util.IOUtils;
 import org.elasticsearch.test.cluster.util.OS;
 import org.elasticsearch.test.cluster.util.Pair;
@@ -649,27 +650,56 @@ private void installPlugins() {
                     .toList();
 
                 List<String> toInstall = spec.getPlugins()
+                    .entrySet()
                     .stream()
                     .map(
-                        pluginName -> pluginPaths.stream()
+                        plugin -> pluginPaths.stream()
                             .map(path -> Pair.of(pattern.matcher(path.getFileName().toString()), path))
-                            .filter(pair -> pair.left.matches() && pair.left.group(1).equals(pluginName))
+                            .filter(pair -> pair.left.matches() && pair.left.group(1).equals(plugin.getKey()))
                             .map(p -> p.right.getParent().resolve(p.left.group(0)))
                             .findFirst()
+                            .map(path -> {
+                                DefaultPluginInstallSpec installSpec = plugin.getValue();
+                                // Path the plugin archive with configured overrides if necessary
+                                if (installSpec.entitlementsOverride != null || installSpec.propertiesOverride != null) {
+                                    Path target;
+                                    try {
+                                        target = Files.createTempFile("patched-", path.getFileName().toString());
+                                    } catch (IOException e) {
+                                        throw new UncheckedIOException("Failed to create temporary file", e);
+                                    }
+                                    ArchivePatcher patcher = new ArchivePatcher(path, target);
+                                    if (installSpec.entitlementsOverride != null) {
+                                        patcher.override(
+                                            "entitlement-policy.yaml",
+                                            original -> installSpec.entitlementsOverride.apply(original).asStream()
+                                        );
+                                    }
+                                    if (installSpec.propertiesOverride != null) {
+                                        patcher.override(
+                                            "plugin-descriptor.properties",
+                                            original -> installSpec.propertiesOverride.apply(original).asStream()
+                                        );
+                                    }
+                                    return patcher.patch();
+                                } else {
+                                    return path;
+                                }
+                            })
                             .orElseThrow(() -> {
                                 String taskPath = System.getProperty("tests.task");
                                 String project = taskPath.substring(0, taskPath.lastIndexOf(':'));
 
-                                throw new RuntimeException(
+                                return new RuntimeException(
                                     "Unable to locate plugin '"
-                                        + pluginName
+                                        + plugin.getKey()
                                         + "'. Ensure you've added the following to the build script for project '"
                                         + project
                                         + "':\n\n"
                                         + "dependencies {\n"
                                         + "  clusterPlugins "
                                         + "project(':plugins:"
-                                        + pluginName
+                                        + plugin.getKey()
                                         + "')"
                                         + "\n}"
                                 );

test/test-clusters/src/main/java/org/elasticsearch/test/cluster/local/AbstractLocalSpecBuilder.java

@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -34,7 +35,7 @@ public abstract class AbstractLocalSpecBuilder<T extends LocalSpecBuilder<?>> im
     private final List<EnvironmentProvider> environmentProviders = new ArrayList<>();
     private final Map<String, String> environment = new HashMap<>();
     private final Set<String> modules = new HashSet<>();
-    private final Set<String> plugins = new HashSet<>();
+    private final Map<String, DefaultPluginInstallSpec> plugins = new HashMap<>();
     private final Set<FeatureFlag> features = EnumSet.noneOf(FeatureFlag.class);
     private final List<SettingsProvider> keystoreProviders = new ArrayList<>();
     private final Map<String, String> keystoreSettings = new HashMap<>();
@@ -132,11 +133,19 @@ Set<String> getModules() {
 
     @Override
     public T plugin(String pluginName) {
-        this.plugins.add(pluginName);
+        this.plugins.put(pluginName, new DefaultPluginInstallSpec());
         return cast(this);
     }
 
-    Set<String> getPlugins() {
+    @Override
+    public T plugin(String pluginName, Consumer<? super PluginInstallSpec> config) {
+        DefaultPluginInstallSpec spec = new DefaultPluginInstallSpec();
+        config.accept(spec);
+        this.plugins.put(pluginName, spec);
+        return cast(this);
+    }
+
+    Map<String, DefaultPluginInstallSpec> getPlugins() {
         return inherit(() -> parent.getPlugins(), plugins);
     }
 

test/test-clusters/src/main/java/org/elasticsearch/test/cluster/local/DefaultPluginInstallSpec.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.test.cluster.local;
+
+import org.elasticsearch.test.cluster.util.resource.Resource;
+
+import java.util.function.Function;
+
+public class DefaultPluginInstallSpec implements PluginInstallSpec {
+    Function<? super String, ? extends Resource> propertiesOverride;
+    Function<? super String, ? extends Resource> entitlementsOverride;
+
+    @Override
+    public PluginInstallSpec withPropertiesOverride(Function<? super String, ? extends Resource> override) {
+        this.propertiesOverride = override;
+        return this;
+    }
+
+    @Override
+    public PluginInstallSpec withEntitlementsOverride(Function<? super String, ? extends Resource> override) {
+        this.entitlementsOverride = override;
+        return this;
+    }
+}

test/test-clusters/src/main/java/org/elasticsearch/test/cluster/local/LocalClusterSpec.java

@@ -92,7 +92,7 @@ public static class LocalNodeSpec {
         private final List<EnvironmentProvider> environmentProviders;
         private final Map<String, String> environment;
         private final Set<String> modules;
-        private final Set<String> plugins;
+        private final Map<String, DefaultPluginInstallSpec> plugins;
         private final DistributionType distributionType;
         private final Set<FeatureFlag> features;
         private final List<SettingsProvider> keystoreProviders;
@@ -114,7 +114,7 @@ public LocalNodeSpec(
             List<EnvironmentProvider> environmentProviders,
             Map<String, String> environment,
             Set<String> modules,
-            Set<String> plugins,
+            Map<String, DefaultPluginInstallSpec> plugins,
             DistributionType distributionType,
             Set<FeatureFlag> features,
             List<SettingsProvider> keystoreProviders,
@@ -179,7 +179,7 @@ public Set<String> getModules() {
             return modules;
         }
 
-        public Set<String> getPlugins() {
+        public Map<String, DefaultPluginInstallSpec> getPlugins() {
             return plugins;
         }
 

test/test-clusters/src/main/java/org/elasticsearch/test/cluster/local/LocalSpecBuilder.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.test.cluster.util.Version;
 import org.elasticsearch.test.cluster.util.resource.Resource;
 
+import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -73,6 +74,11 @@ interface LocalSpecBuilder<T extends LocalSpecBuilder<?>> {
      */
     T plugin(String pluginName);
 
+    /**
+     * Ensure plugin is installed into the distribution.
+     */
+    T plugin(String pluginName, Consumer<? super PluginInstallSpec> config);
+
     /**
      * Require feature to be enabled in the cluster.
      */

test/test-clusters/src/main/java/org/elasticsearch/test/cluster/local/PluginInstallSpec.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.test.cluster.local;
+
+import org.elasticsearch.test.cluster.util.resource.Resource;
+
+import java.util.function.Function;
+
+public interface PluginInstallSpec {
+
+    /**
+     * Override bundled plugin properties file with the given {@link Resource}. The provided override function receives the original
+     * file content as function argument.
+     *
+     * @param override function returning resource used to override bundled properties file
+     */
+    PluginInstallSpec withPropertiesOverride(Function<? super String, ? extends Resource> override);
+
+    /**
+     * Override bundled entitlements policy file with the given {@link Resource}. The provided override function receives the original
+     * file content as function argument.
+     *
+     * @param override function returning resource used to override bundled entitlements policy file
+     */
+    PluginInstallSpec withEntitlementsOverride(Function<? super String, ? extends Resource> override);
+}

test/test-clusters/src/main/java/org/elasticsearch/test/cluster/util/ArchivePatcher.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.test.cluster.util;
+
+import java.io.BufferedOutputStream;
+import java.io.BufferedReader;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.UncheckedIOException;
+import java.nio.file.Path;
+import java.util.Enumeration;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipFile;
+import java.util.zip.ZipOutputStream;
+
+public class ArchivePatcher {
+    private final Path original;
+    private final Path target;
+    private final Map<String, Function<? super String, ? extends InputStream>> overrides = new HashMap<>();
+
+    public ArchivePatcher(Path original, Path target) {
+        this.original = original;
+        this.target = target;
+    }
+
+    public void override(String filename, Function<? super String, ? extends InputStream> override) {
+        this.overrides.put(filename, override);
+    }
+
+    public Path patch() {
+        try (
+            ZipFile input = new ZipFile(original.toFile());
+            ZipOutputStream output = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(target.toFile())))
+        ) {
+            Enumeration<? extends ZipEntry> entries = input.entries();
+            while (entries.hasMoreElements()) {
+                ZipEntry entry = entries.nextElement();
+                output.putNextEntry(entry);
+                if (overrides.containsKey(entry.getName())) {
+                    try (BufferedReader reader = new BufferedReader(new InputStreamReader(input.getInputStream(entry)))) {
+                        String content = reader.lines().collect(Collectors.joining(System.lineSeparator()));
+                        overrides.get(entry.getName()).apply(content).transferTo(output);
+                    }
+                } else {
+                    input.getInputStream(entry).transferTo(output);
+                }
+                output.closeEntry();
+            }
+            output.flush();
+            output.finish();
+        } catch (IOException e) {
+            throw new UncheckedIOException("Failed to patch archive", e);
+        }
+
+        return target;
+    }
+}