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:
Ryan Ernst 2018-09-27 12:23:59 -07:00 committed by GitHub
parent 73ee721b29
commit a2c941806b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
15 changed files with 52 additions and 34 deletions

View File

@ -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)

View File

@ -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();

View File

@ -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)

View File

@ -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");
}

View File

@ -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);
}

View File

@ -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,

View File

@ -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());

View File

@ -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;
}

View File

@ -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 + "]");
}

View File

@ -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;
}

View File

@ -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());
}
}));
}

View File

@ -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() {

View File

@ -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));

View File

@ -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);

View File

@ -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);