diff --git a/plugin/src/test/java/org/elasticsearch/script/MockMustacheScriptEngine.java b/plugin/src/test/java/org/elasticsearch/script/MockMustacheScriptEngine.java
index 57df1936053..42fb5a249d9 100644
--- a/plugin/src/test/java/org/elasticsearch/script/MockMustacheScriptEngine.java
+++ b/plugin/src/test/java/org/elasticsearch/script/MockMustacheScriptEngine.java
@@ -39,11 +39,13 @@ public class MockMustacheScriptEngine extends MockScriptEngine {
     }
 
     @Override
-    public Object compile(String name, String script, Map<String, String> params) {
+    public <T> T compile(String name, String script, ScriptContext<T> context, Map<String, String> params) {
         if (script.contains("{{") && script.contains("}}")) {
             throw new IllegalArgumentException("Fix your test to not rely on mustache");
         }
-        // We always return the script's source as it is
-        return new MockCompiledScript(name, params, script, null);
+        if (context.instanceClazz.equals(ExecutableScript.class) == false) {
+            throw new IllegalArgumentException("mock mustache only understands template scripts, not [" + context.name + "]");
+        }
+        return context.compiledClazz.cast((ExecutableScript.Compiled) vars -> new MockExecutableScript(vars, p -> script));
     }
 }
diff --git a/plugin/src/test/java/org/elasticsearch/xpack/monitoring/test/MockPainlessScriptEngine.java b/plugin/src/test/java/org/elasticsearch/xpack/monitoring/test/MockPainlessScriptEngine.java
index 5bbc5e16c4e..bf90f89989b 100644
--- a/plugin/src/test/java/org/elasticsearch/xpack/monitoring/test/MockPainlessScriptEngine.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/monitoring/test/MockPainlessScriptEngine.java
@@ -6,9 +6,12 @@
 package org.elasticsearch.xpack.monitoring.test;
 
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.MockScriptEngine;
 import org.elasticsearch.script.MockScriptPlugin;
+import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
+import org.elasticsearch.script.SearchScript;
 
 import java.util.Collections;
 import java.util.Map;
@@ -39,8 +42,12 @@ public class MockPainlessScriptEngine extends MockScriptEngine {
     }
 
     @Override
-    public Object compile(String name, String script, Map<String, String> params) {
-        // We always return the script's source as it is
-        return new MockCompiledScript(name, params, script, null);
+    public <T> T compile(String name, String script, ScriptContext<T> context, Map<String, String> params) {
+        if (context.instanceClazz.equals(ExecutableScript.class)) {
+            return context.compiledClazz.cast((ExecutableScript.Compiled) vars -> new MockExecutableScript(vars, p -> script));
+        } else if (context.instanceClazz.equals(SearchScript.class)) {
+            return context.compiledClazz.cast((SearchScript.Compiled) (vars, lookup) -> new MockSearchScript(lookup, vars, p -> script));
+        }
+        throw new IllegalArgumentException("mock painless does not know how to handle context [" + context.name + "]");
     }
 }