Tests: Add support for custom contexts to mock scripts (#34100)

This commit adds the ability to plug in compilation of custom contexts
in mock script engine. This is needed for testing plugins which add
custom contexts like watcher.

modules/ingest-common/src/test/java/org/elasticsearch/ingest/common/ScriptProcessorTests.java

@@ -19,10 +19,6 @@
 
 package org.elasticsearch.ingest.common;
 
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.ingest.IngestDocument;
 import org.elasticsearch.ingest.RandomDocumentPicks;
@@ -33,6 +29,10 @@ import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.test.ESTestCase;
 
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
 import static org.hamcrest.Matchers.hasKey;
 import static org.hamcrest.core.Is.is;
 
@@ -52,7 +52,8 @@ public class ScriptProcessorTests extends ESTestCase {
                             ctx.put("bytes_total", randomBytesTotal);
                             return null;
                         }
-                    )
+                    ),
+                    Collections.emptyMap()
                 )
             ),
             new HashMap<>(ScriptModule.CORE_CONTEXTS)

server/src/test/java/org/elasticsearch/action/update/UpdateRequestTests.java

@@ -132,7 +132,7 @@ public class UpdateRequestTests extends ESTestCase {
                     return null;
                 });
         scripts.put("return", vars -> null);
-        final MockScriptEngine engine = new MockScriptEngine("mock", scripts);
+        final MockScriptEngine engine = new MockScriptEngine("mock", scripts, Collections.emptyMap());
         Map<String, ScriptEngine> engines = Collections.singletonMap(engine.getType(), engine);
         ScriptService scriptService = new ScriptService(baseSettings, engines, ScriptModule.CORE_CONTEXTS);
         final Settings settings = settings(Version.CURRENT).build();

server/src/test/java/org/elasticsearch/ingest/ConditionalProcessorTests.java

@@ -19,13 +19,6 @@
 
 package org.elasticsearch.ingest;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-import java.util.function.Consumer;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.script.MockScriptEngine;
 import org.elasticsearch.script.Script;
@@ -34,6 +27,14 @@ import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.test.ESTestCase;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
+
 import static org.hamcrest.Matchers.hasKey;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.not;
@@ -52,7 +53,8 @@ public class ConditionalProcessorTests extends ESTestCase {
                     Script.DEFAULT_SCRIPT_LANG,
                     Collections.singletonMap(
                         scriptName, ctx -> trueValue.equals(ctx.get(conditionalField))
-                    )
+                    ),
+                    Collections.emptyMap()
                 )
             ),
             new HashMap<>(ScriptModule.CORE_CONTEXTS)
@@ -120,7 +122,8 @@ public class ConditionalProcessorTests extends ESTestCase {
                             }
                             return false;
                         }
-                    )
+                    ),
+                    Collections.emptyMap()
                 )
             ),
             new HashMap<>(ScriptModule.CORE_CONTEXTS)

server/src/test/java/org/elasticsearch/script/ScriptServiceTests.java

@@ -66,12 +66,12 @@ public class ScriptServiceTests extends ESTestCase {
             scripts.put(i + "+" + i, p -> null); // only care about compilation, not execution
         }
         scripts.put("script", p -> null);
-        scriptEngine = new MockScriptEngine(Script.DEFAULT_SCRIPT_LANG, scripts);
+        scriptEngine = new MockScriptEngine(Script.DEFAULT_SCRIPT_LANG, scripts, Collections.emptyMap());
         //prevent duplicates using map
         contexts = new HashMap<>(ScriptModule.CORE_CONTEXTS);
         engines = new HashMap<>();
         engines.put(scriptEngine.getType(), scriptEngine);
-        engines.put("test", new MockScriptEngine("test", scripts));
+        engines.put("test", new MockScriptEngine("test", scripts, Collections.emptyMap()));
         logger.info("--> setup script service");
     }
 

server/src/test/java/org/elasticsearch/search/aggregations/metrics/InternalScriptedMetricTests.java

@@ -30,8 +30,6 @@ import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.search.aggregations.Aggregation.CommonFields;
 import org.elasticsearch.search.aggregations.ParsedAggregation;
-import org.elasticsearch.search.aggregations.metrics.InternalScriptedMetric;
-import org.elasticsearch.search.aggregations.metrics.ParsedScriptedMetric;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.test.InternalAggregationTestCase;
 
@@ -118,7 +116,8 @@ public class InternalScriptedMetricTests extends InternalAggregationTestCase<Int
         // mock script always retuns the size of the input aggs list as result
         @SuppressWarnings("unchecked")
         MockScriptEngine scriptEngine = new MockScriptEngine(MockScriptEngine.NAME,
-                Collections.singletonMap(REDUCE_SCRIPT_NAME, script -> ((List<Object>) script.get("states")).size()));
+                Collections.singletonMap(REDUCE_SCRIPT_NAME, script -> ((List<Object>) script.get("states")).size()),
+                Collections.emptyMap());
         Map<String, ScriptEngine> engines = Collections.singletonMap(scriptEngine.getType(), scriptEngine);
         return new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS);
     }

server/src/test/java/org/elasticsearch/search/aggregations/metrics/ScriptedMetricAggregatorTests.java

@@ -35,8 +35,6 @@ import org.elasticsearch.script.ScriptModule;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.search.aggregations.AggregatorTestCase;
-import org.elasticsearch.search.aggregations.metrics.ScriptedMetric;
-import org.elasticsearch.search.aggregations.metrics.ScriptedMetricAggregationBuilder;
 import org.junit.BeforeClass;
 
 import java.io.IOException;
@@ -345,7 +343,7 @@ public class ScriptedMetricAggregatorTests extends AggregatorTestCase {
      */
     @Override
     protected QueryShardContext queryShardContextMock(MapperService mapperService) {
-        MockScriptEngine scriptEngine = new MockScriptEngine(MockScriptEngine.NAME, SCRIPTS);
+        MockScriptEngine scriptEngine = new MockScriptEngine(MockScriptEngine.NAME, SCRIPTS, Collections.emptyMap());
         Map<String, ScriptEngine> engines = Collections.singletonMap(scriptEngine.getType(), scriptEngine);
         ScriptService scriptService =  new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS);
         return new QueryShardContext(0, mapperService.getIndexSettings(), null, null, mapperService, null, scriptService,

server/src/test/java/org/elasticsearch/search/sort/AbstractSortTestCase.java

@@ -84,7 +84,7 @@ public abstract class AbstractSortTestCase<T extends SortBuilder<T>> extends EST
                 .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir().toString())
                 .build();
         Map<String, Function<Map<String, Object>, Object>> scripts = Collections.singletonMap(MOCK_SCRIPT_NAME, p -> null);
-        ScriptEngine engine = new MockScriptEngine(MockScriptEngine.NAME, scripts);
+        ScriptEngine engine = new MockScriptEngine(MockScriptEngine.NAME, scripts, Collections.emptyMap());
         scriptService = new ScriptService(baseSettings, Collections.singletonMap(engine.getType(), engine), ScriptModule.CORE_CONTEXTS);
 
         SearchModule searchModule = new SearchModule(Settings.EMPTY, false, emptyList());

test/framework/src/main/java/org/elasticsearch/ingest/TestTemplateService.java

@@ -43,8 +43,7 @@ public class TestTemplateService extends ScriptService {
     }
 
     private TestTemplateService(boolean compilationException) {
-        super(Settings.EMPTY, Collections.singletonMap(DEFAULT_TEMPLATE_LANG,
+        super(Settings.EMPTY, Collections.singletonMap(DEFAULT_TEMPLATE_LANG, new MockScriptEngine()), Collections.emptyMap());
-            new MockScriptEngine(MockScriptEngine.NAME, Collections.emptyMap())), Collections.emptyMap());
         this.compilationException = compilationException;
     }
 

test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java

@@ -53,18 +53,26 @@ import static java.util.Collections.emptyMap;
  */
 public class MockScriptEngine implements ScriptEngine {
 
+    /** A non-typed compiler for a single custom context */
+    public interface ContextCompiler {
+        Object compile(Function<Map<String, Object>, Object> script, Map<String, String> params);
+    }
+
     public static final String NAME = "mockscript";
 
     private final String type;
     private final Map<String, Function<Map<String, Object>, Object>> scripts;
+    private final Map<ScriptContext<?>, ContextCompiler> contexts;
 
-    public MockScriptEngine(String type, Map<String, Function<Map<String, Object>, Object>> scripts) {
+    public MockScriptEngine(String type, Map<String, Function<Map<String, Object>, Object>> scripts,
+                            Map<ScriptContext<?>, ContextCompiler> contexts) {
         this.type = type;
         this.scripts = Collections.unmodifiableMap(scripts);
+        this.contexts = Collections.unmodifiableMap(contexts);
     }
 
     public MockScriptEngine() {
-        this(NAME, Collections.emptyMap());
+        this(NAME, Collections.emptyMap(), Collections.emptyMap());
     }
 
     @Override
@@ -198,6 +206,10 @@ public class MockScriptEngine implements ScriptEngine {
             ScriptedMetricAggContexts.ReduceScript.Factory factory = mockCompiled::createMetricAggReduceScript;
             return context.factoryClazz.cast(factory);
         }
+        ContextCompiler compiler = contexts.get(context);
+        if (compiler != null) {
+            return context.factoryClazz.cast(compiler.compile(script, params));
+        }
         throw new IllegalArgumentException("mock script engine does not know how to handle context [" + context.name + "]");
     }
 

test/framework/src/main/java/org/elasticsearch/script/MockScriptPlugin.java

@@ -24,6 +24,7 @@ import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.ScriptPlugin;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Map;
 import java.util.function.Function;
 
@@ -36,11 +37,15 @@ public abstract class MockScriptPlugin extends Plugin implements ScriptPlugin {
 
     @Override
     public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
-        return new MockScriptEngine(pluginScriptLang(), pluginScripts());
+        return new MockScriptEngine(pluginScriptLang(), pluginScripts(), pluginContextCompilers());
     }
 
     protected abstract Map<String, Function<Map<String, Object>, Object>> pluginScripts();
 
+    protected Map<ScriptContext<?>, MockScriptEngine.ContextCompiler> pluginContextCompilers() {
+        return Collections.emptyMap();
+    }
+
     public String pluginScriptLang() {
         return NAME;
     }

test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java

@@ -1384,7 +1384,7 @@ public abstract class ESTestCase extends LuceneTestCase {
         return new ScriptModule(Settings.EMPTY, singletonList(new ScriptPlugin() {
             @Override
             public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
-                return new MockScriptEngine(MockScriptEngine.NAME, Collections.singletonMap("1", script -> "1"));
+                return new MockScriptEngine(MockScriptEngine.NAME, Collections.singletonMap("1", script -> "1"), Collections.emptyMap());
             }
         }));
     }

x-pack/plugin/upgrade/src/test/java/org/elasticsearch/xpack/upgrade/IndexUpgradeTasksIT.java

@@ -88,7 +88,7 @@ public class IndexUpgradeTasksIT extends ESIntegTestCase {
 
         @Override
         public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
-            return new MockScriptEngine(pluginScriptLang(), pluginScripts());
+            return new MockScriptEngine(pluginScriptLang(), pluginScripts(), Collections.emptyMap());
         }
 
         public String pluginScriptLang() {

x-pack/plugin/watcher/src/test/java/org/elasticsearch/xpack/watcher/condition/ScriptConditionTests.java

@@ -83,7 +83,7 @@ public class ScriptConditionTests extends ESTestCase {
             return total > threshold;
         });
 
-        ScriptEngine engine = new MockScriptEngine(MockScriptEngine.NAME, scripts);
+        ScriptEngine engine = new MockScriptEngine(MockScriptEngine.NAME, scripts, Collections.emptyMap());
         scriptService = new ScriptService(Settings.EMPTY, Collections.singletonMap(engine.getType(), engine),
             Collections.singletonMap(Watcher.SCRIPT_EXECUTABLE_CONTEXT.name, Watcher.SCRIPT_EXECUTABLE_CONTEXT));
 

x-pack/plugin/watcher/src/test/java/org/elasticsearch/xpack/watcher/input/transform/TransformInputTests.java

@@ -47,7 +47,8 @@ public class TransformInputTests extends ESTestCase {
     @Before
     public void setupScriptService() {
         Map<String, ScriptEngine> engines = new HashMap<>();
-        engines.put(MockScriptEngine.NAME, new MockScriptEngine(MockScriptEngine.NAME, Collections.singletonMap("1", s -> "2")));
+        engines.put(MockScriptEngine.NAME,
+            new MockScriptEngine(MockScriptEngine.NAME, Collections.singletonMap("1", s -> "2"), Collections.emptyMap()));
         Map<String, ScriptContext<?>> contexts = new HashMap<>();
         contexts.put(Watcher.SCRIPT_TEMPLATE_CONTEXT.name, Watcher.SCRIPT_TEMPLATE_CONTEXT);
         contexts.put(Watcher.SCRIPT_SEARCH_CONTEXT.name, Watcher.SCRIPT_SEARCH_CONTEXT);

x-pack/plugin/watcher/src/test/java/org/elasticsearch/xpack/watcher/test/integration/SearchTransformTests.java

@@ -70,7 +70,7 @@ public class SearchTransformTests extends ESTestCase {
         XContentParser parser = createParser(builder);
         parser.nextToken();
 
-        final MockScriptEngine engine = new MockScriptEngine("mock", Collections.emptyMap());
+        final MockScriptEngine engine = new MockScriptEngine("mock", Collections.emptyMap(), Collections.emptyMap());
         Map<String, ScriptEngine> engines = Collections.singletonMap(engine.getType(), engine);
         ScriptService scriptService = new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS);