Test: MockScoreScript can be cacheable (#55422)

Author: stu-elastic

server/src/test/java/org/elasticsearch/index/query/ScriptScoreQueryBuilderTests.java

@@ -94,7 +94,7 @@ public void testIllegalArguments() {
     }
 
     /**
-     * Check that this query is generally not cacheable
+     * Check that this query is cacheable
      */
     @Override
     public void testCacheability() throws IOException {
@@ -105,7 +105,7 @@ public void testCacheability() throws IOException {
         QueryShardContext context = createShardContext();
         QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
         assertNotNull(rewriteQuery.toQuery(context));
-        assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
+        assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 
     @Override

server/src/test/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java

@@ -837,29 +837,28 @@ public void testCacheability() throws IOException {
         assertEquals("query should " + (isCacheable ? "" : "not") + " be cacheable: " + queryBuilder.toString(), isCacheable,
                 context.isCacheable());
 
-        // check the two non-cacheable cases explicitly
         ScoreFunctionBuilder<?> scriptScoreFunction = new ScriptScoreFunctionBuilder(
                 new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap()));
-        RandomScoreFunctionBuilder randomScoreFunctionBuilder = new RandomScoreFunctionBuilderWithFixedSeed();
-
-        for (ScoreFunctionBuilder<?> scoreFunction : List.of(scriptScoreFunction, randomScoreFunctionBuilder)) {
-            FilterFunctionBuilder[] functions = new FilterFunctionBuilder[] {
-                    new FilterFunctionBuilder(RandomQueryBuilder.createQuery(random()), scoreFunction) };
-            queryBuilder = new FunctionScoreQueryBuilder(functions);
+        queryBuilder = new FunctionScoreQueryBuilder(new FilterFunctionBuilder[] {
+            new FilterFunctionBuilder(RandomQueryBuilder.createQuery(random()), scriptScoreFunction) });
+        context = createShardContext();
+        rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertTrue("function script query should be cacheable" + queryBuilder.toString(), context.isCacheable());
 
-            context = createShardContext();
-            rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
-            assertNotNull(rewriteQuery.toQuery(context));
-            assertFalse("query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
-        }
+        RandomScoreFunctionBuilder randomScoreFunctionBuilder = new RandomScoreFunctionBuilderWithFixedSeed();
+        queryBuilder = new FunctionScoreQueryBuilder(new FilterFunctionBuilder[] {
+            new FilterFunctionBuilder(RandomQueryBuilder.createQuery(random()), randomScoreFunctionBuilder) });
+        context = createShardContext();
+        rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
+        assertNotNull(rewriteQuery.toQuery(context));
+        assertFalse("function random query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
     }
 
     private boolean isCacheable(FunctionScoreQueryBuilder queryBuilder) {
         FilterFunctionBuilder[] filterFunctionBuilders = queryBuilder.filterFunctionBuilders();
         for (FilterFunctionBuilder builder : filterFunctionBuilders) {
-            if (builder.getScoreFunction() instanceof ScriptScoreFunctionBuilder) {
-                return false;
-            } else if (builder.getScoreFunction() instanceof RandomScoreFunctionBuilder
+            if (builder.getScoreFunction() instanceof RandomScoreFunctionBuilder
                 && ((RandomScoreFunctionBuilder) builder.getScoreFunction()).getSeed() == null) {
                 return false;
             }

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

@@ -225,7 +225,7 @@ public double execute(Map<String, Object> params1, double[] values) {
             };
             return context.factoryClazz.cast(factory);
         } else if (context.instanceClazz.equals(ScoreScript.class)) {
-            ScoreScript.Factory factory = new MockScoreScript(script::apply);
+            ScoreScript.Factory factory = new MockScoreScript(script);
             return context.factoryClazz.cast(factory);
         } else if (context.instanceClazz.equals(ScriptedMetricAggContexts.InitScript.class)) {
             ScriptedMetricAggContexts.InitScript.Factory factory = new MockMetricAggInitScriptFactory(script);
@@ -545,9 +545,9 @@ public static Script mockInlineScript(final String script) {
 
     public class MockScoreScript implements ScoreScript.Factory {
 
-        private final Function<Map<String, Object>, Object> script;
+        private final MockDeterministicScript script;
 
-        public MockScoreScript(Function<Map<String, Object>, Object> script) {
+        public MockScoreScript(MockDeterministicScript script) {
             this.script = script;
         }
 
@@ -581,6 +581,11 @@ public void setScorer(Scorable scorer) {
                 }
             };
         }
+
+        @Override
+        public boolean isResultDeterministic() {
+            return script.isResultDeterministic();
+        }
     }
 
     class MockAggregationScript implements AggregationScript.Factory {

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/monitoring/test/MockPainlessScriptEngine.java

@@ -6,6 +6,7 @@
 package org.elasticsearch.xpack.core.monitoring.test;
 
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.script.MockDeterministicScript;
 import org.elasticsearch.script.MockScriptEngine;
 import org.elasticsearch.script.MockScriptPlugin;
 import org.elasticsearch.script.ScoreScript;
@@ -44,7 +45,7 @@ public String getType() {
     @Override
     public <T> T compile(String name, String script, ScriptContext<T> context, Map<String, String> options) {
         if (context.instanceClazz.equals(ScoreScript.class)) {
-            return context.factoryClazz.cast(new MockScoreScript(p -> 0.0));
+            return context.factoryClazz.cast(new MockScoreScript(MockDeterministicScript.asDeterministic(p -> 0.0)));
         }
         throw new IllegalArgumentException("mock painless does not know how to handle context [" + context.name + "]");
     }

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeIntegTestCase.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.persistent.PersistentTaskState;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.script.IngestScript;
+import org.elasticsearch.script.MockDeterministicScript;
 import org.elasticsearch.script.MockScriptEngine;
 import org.elasticsearch.script.MockScriptPlugin;
 import org.elasticsearch.script.ScoreScript;
@@ -307,7 +308,7 @@ public String getType() {
         @Override
         public <T> T compile(String name, String script, ScriptContext<T> context, Map<String, String> options) {
             if (context.instanceClazz.equals(ScoreScript.class)) {
-                return context.factoryClazz.cast(new MockScoreScript(p -> 0.0));
+                return context.factoryClazz.cast(new MockScoreScript(MockDeterministicScript.asDeterministic(p -> 0.0)));
             }
             if (context.name.equals("ingest")) {
                 IngestScript.Factory factory = vars -> new IngestScript(vars) {