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.
This commit is contained in:
parent
73ee721b29
commit
a2c941806b
|
@ -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)
|
||||
|
|
|
@ -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();
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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");
|
||||
}
|
||||
|
||||
|
|
|
@ -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);
|
||||
}
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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());
|
||||
|
|
|
@ -43,8 +43,7 @@ public class TestTemplateService extends ScriptService {
|
|||
}
|
||||
|
||||
private TestTemplateService(boolean compilationException) {
|
||||
super(Settings.EMPTY, Collections.singletonMap(DEFAULT_TEMPLATE_LANG,
|
||||
new MockScriptEngine(MockScriptEngine.NAME, Collections.emptyMap())), Collections.emptyMap());
|
||||
super(Settings.EMPTY, Collections.singletonMap(DEFAULT_TEMPLATE_LANG, new MockScriptEngine()), Collections.emptyMap());
|
||||
this.compilationException = compilationException;
|
||||
}
|
||||
|
||||
|
|
|
@ -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 + "]");
|
||||
}
|
||||
|
||||
|
|
|
@ -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;
|
||||
}
|
||||
|
|
|
@ -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());
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
|
|
@ -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() {
|
||||
|
|
|
@ -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));
|
||||
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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);
|
||||
|
||||
|
|
Loading…
Reference in New Issue