interpretation = executor.submit(() ->
+ driver.interpret(interpreter, "block-until-cancelled", runningContext));
+
+ try {
+ assertTrue(
+ driver.probe().awaitInterpretStarted(ASYNC_TIMEOUT),
+ "Interpreter did not start");
+ assertEquals(JobStatus.RUNNING, driver.getStatus(interpreter, "running-paragraph"));
+ assertEquals(42, driver.getProgress(interpreter, runningContext));
+
+ driver.cancel(interpreter, runningContext);
+ assertTrue(
+ driver.probe().awaitCancelObserved(ASYNC_TIMEOUT),
+ "Interpreter did not observe cancellation");
+
+ ProbeSnapshot snapshot = driver.probe().snapshot();
+ assertEquals("contract-note", snapshot.getCancelNoteId());
+ assertEquals("running-paragraph", snapshot.getCancelParagraphId());
+
+ ContractResult terminalResult =
+ interpretation.get(ASYNC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
+ assertEquals(ResultCode.SUCCESS, terminalResult.getCode());
+ assertEquals(1, terminalResult.getMessages().size());
+ assertEquals(ResultType.TEXT, terminalResult.getMessages().get(0).getType());
+ assertEquals("cancelled", terminalResult.getMessages().get(0).getData());
+ } finally {
+ driver.probe().releaseInterpretation();
+ if (!interpretation.isDone()) {
+ driver.faults().abortPendingCalls();
+ }
+ interpretation.cancel(true);
+ executor.shutdownNow();
+ assertTrue(
+ executor.awaitTermination(ASYNC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS),
+ "Interpret executor did not terminate");
+ }
+ }
+
+ @Test
+ protected void shouldNormalizeExecutionOperationInterpreterAndTransportFailures()
+ throws Exception {
+ InterpreterRef interpreter = interpreterSpec.getInterpreter();
+
+ ContractResult result =
+ driver.interpret(interpreter, "throw-application-error", context("application-error"));
+ assertEquals(ResultCode.ERROR, result.getCode());
+ assertTrue(result.getMessages().get(0).getData().contains("contract interpret failure"));
+
+ assertFailureCategory(
+ FailureCategory.OPERATION_FAILED,
+ () -> driver.completion(
+ interpreter,
+ "throw-completion-error",
+ 0,
+ context("completion-error")));
+
+ InterpreterRef missingInterpreter = interpreter.withSessionId("missing-session");
+ assertFailureCategory(
+ FailureCategory.INTERPRETER_NOT_FOUND,
+ () -> driver.getFormType(missingInterpreter));
+
+ driver.faults().makeTransportUnavailable();
+ assertFailureCategory(
+ FailureCategory.TRANSPORT_UNAVAILABLE,
+ () -> driver.getStatus(interpreter, "missing-job"));
+ }
+
+ private ContractContext context(String paragraphId) {
+ return new ContractContext(
+ "contract-note",
+ "Contract Note",
+ paragraphId,
+ "contract",
+ "Contract Paragraph",
+ "contract text",
+ "contract-user",
+ Set.of("contract-role", "contract-auditor"),
+ "contract-ticket",
+ Map.of("contract-local", "local-value"),
+ Map.of("request-config", "config-value"),
+ Map.of("request-gui", "gui-value"),
+ Map.of("request-note-gui", "note-gui-value"));
+ }
+
+ private void assertFailureCategory(
+ FailureCategory expectedCategory, ThrowingOperation operation) {
+ ContractFailure failure = assertThrows(ContractFailure.class, operation::run);
+ assertEquals(expectedCategory, failure.getCategory());
+ }
+
+ @FunctionalInterface
+ private interface ThrowingOperation {
+ void run() throws ContractFailure;
+ }
+}
diff --git a/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/InterpreterRpcContractDriver.java b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/InterpreterRpcContractDriver.java
new file mode 100644
index 00000000000..a03d5807575
--- /dev/null
+++ b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/InterpreterRpcContractDriver.java
@@ -0,0 +1,487 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.zeppelin.interpreter.remote;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * Transport-neutral test contract for Server-to-Interpreter control RPC implementations.
+ *
+ * Transport adapters normalize their generated messages and failures into the value types in
+ * this interface. The shared scenarios can therefore be reused without depending on a wire
+ * protocol, client topology, or generated RPC classes.
+ */
+public interface InterpreterRpcContractDriver extends AutoCloseable {
+
+ InterpreterSpec interpreterSpec();
+
+ /** Starts the transport and performs its process-level initialization. */
+ void start() throws Exception;
+
+ void createInterpreter(InterpreterSpec spec) throws ContractFailure;
+
+ FormType getFormType(InterpreterRef interpreter) throws ContractFailure;
+
+ int getProgress(InterpreterRef interpreter, ContractContext context) throws ContractFailure;
+
+ List completion(
+ InterpreterRef interpreter, String buffer, int cursor, ContractContext context)
+ throws ContractFailure;
+
+ ContractResult interpret(
+ InterpreterRef interpreter, String statement, ContractContext context)
+ throws ContractFailure;
+
+ JobStatus getStatus(InterpreterRef interpreter, String jobId) throws ContractFailure;
+
+ void cancel(InterpreterRef interpreter, ContractContext context) throws ContractFailure;
+
+ void closeInterpreter(InterpreterRef interpreter) throws ContractFailure;
+
+ ContractProbe probe();
+
+ ContractFaults faults();
+
+ @Override
+ void close() throws Exception;
+
+ enum FormType {
+ NATIVE,
+ SIMPLE,
+ NONE
+ }
+
+ enum JobStatus {
+ UNKNOWN,
+ READY,
+ PENDING,
+ RUNNING,
+ FINISHED,
+ ERROR,
+ ABORT
+ }
+
+ enum ResultCode {
+ SUCCESS,
+ INCOMPLETE,
+ ERROR,
+ KEEP_PREVIOUS_RESULT
+ }
+
+ enum ResultType {
+ TEXT,
+ HTML,
+ ANGULAR,
+ TABLE,
+ IMG,
+ SVG,
+ NULL,
+ NETWORK
+ }
+
+ enum FailureCategory {
+ OPERATION_FAILED,
+ INTERPRETER_NOT_FOUND,
+ TRANSPORT_UNAVAILABLE
+ }
+
+ final class InterpreterRef {
+ private final String interpreterGroupId;
+ private final String sessionId;
+ private final String className;
+
+ public InterpreterRef(String interpreterGroupId, String sessionId, String className) {
+ this.interpreterGroupId = interpreterGroupId;
+ this.sessionId = sessionId;
+ this.className = className;
+ }
+
+ public String getInterpreterGroupId() {
+ return interpreterGroupId;
+ }
+
+ public String getSessionId() {
+ return sessionId;
+ }
+
+ public String getClassName() {
+ return className;
+ }
+
+ public InterpreterRef withSessionId(String replacementSessionId) {
+ return new InterpreterRef(interpreterGroupId, replacementSessionId, className);
+ }
+ }
+
+ final class InterpreterSpec {
+ private final InterpreterRef interpreter;
+ private final String userName;
+ private final Map properties;
+
+ public InterpreterSpec(
+ InterpreterRef interpreter, String userName, Map properties) {
+ this.interpreter = interpreter;
+ this.userName = userName;
+ this.properties = immutableCopy(properties);
+ }
+
+ public InterpreterRef getInterpreter() {
+ return interpreter;
+ }
+
+ public String getUserName() {
+ return userName;
+ }
+
+ public Map getProperties() {
+ return properties;
+ }
+ }
+
+ final class ContractContext {
+ private final String noteId;
+ private final String noteName;
+ private final String paragraphId;
+ private final String replName;
+ private final String paragraphTitle;
+ private final String paragraphText;
+ private final String userName;
+ private final Set userRoles;
+ private final String userTicket;
+ private final Map localProperties;
+ private final Map config;
+ private final Map guiParameters;
+ private final Map noteGuiParameters;
+
+ public ContractContext(
+ String noteId,
+ String noteName,
+ String paragraphId,
+ String replName,
+ String paragraphTitle,
+ String paragraphText,
+ String userName,
+ Set userRoles,
+ String userTicket,
+ Map localProperties,
+ Map config,
+ Map guiParameters,
+ Map noteGuiParameters) {
+ this.noteId = noteId;
+ this.noteName = noteName;
+ this.paragraphId = paragraphId;
+ this.replName = replName;
+ this.paragraphTitle = paragraphTitle;
+ this.paragraphText = paragraphText;
+ this.userName = userName;
+ this.userRoles = Collections.unmodifiableSet(new LinkedHashSet<>(userRoles));
+ this.userTicket = userTicket;
+ this.localProperties = immutableCopy(localProperties);
+ this.config = immutableCopy(config);
+ this.guiParameters = immutableCopy(guiParameters);
+ this.noteGuiParameters = immutableCopy(noteGuiParameters);
+ }
+
+ public String getNoteId() {
+ return noteId;
+ }
+
+ public String getNoteName() {
+ return noteName;
+ }
+
+ public String getParagraphId() {
+ return paragraphId;
+ }
+
+ public String getReplName() {
+ return replName;
+ }
+
+ public String getParagraphTitle() {
+ return paragraphTitle;
+ }
+
+ public String getParagraphText() {
+ return paragraphText;
+ }
+
+ public String getUserName() {
+ return userName;
+ }
+
+ public Set getUserRoles() {
+ return userRoles;
+ }
+
+ public String getUserTicket() {
+ return userTicket;
+ }
+
+ public Map getLocalProperties() {
+ return localProperties;
+ }
+
+ public Map getConfig() {
+ return config;
+ }
+
+ public Map getGuiParameters() {
+ return guiParameters;
+ }
+
+ public Map getNoteGuiParameters() {
+ return noteGuiParameters;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (!(other instanceof ContractContext)) {
+ return false;
+ }
+ ContractContext that = (ContractContext) other;
+ return Objects.equals(noteId, that.noteId)
+ && Objects.equals(noteName, that.noteName)
+ && Objects.equals(paragraphId, that.paragraphId)
+ && Objects.equals(replName, that.replName)
+ && Objects.equals(paragraphTitle, that.paragraphTitle)
+ && Objects.equals(paragraphText, that.paragraphText)
+ && Objects.equals(userName, that.userName)
+ && Objects.equals(userRoles, that.userRoles)
+ && Objects.equals(userTicket, that.userTicket)
+ && Objects.equals(localProperties, that.localProperties)
+ && Objects.equals(config, that.config)
+ && Objects.equals(guiParameters, that.guiParameters)
+ && Objects.equals(noteGuiParameters, that.noteGuiParameters);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(
+ noteId,
+ noteName,
+ paragraphId,
+ replName,
+ paragraphTitle,
+ paragraphText,
+ userName,
+ userRoles,
+ userTicket,
+ localProperties,
+ config,
+ guiParameters,
+ noteGuiParameters);
+ }
+ }
+
+ final class ContractCompletion {
+ private final String name;
+ private final String value;
+ private final String meta;
+
+ public ContractCompletion(String name, String value, String meta) {
+ this.name = name;
+ this.value = value;
+ this.meta = meta;
+ }
+
+ public String getName() {
+ return name;
+ }
+
+ public String getValue() {
+ return value;
+ }
+
+ public String getMeta() {
+ return meta;
+ }
+ }
+
+ final class ContractResultMessage {
+ private final ResultType type;
+ private final String data;
+
+ public ContractResultMessage(ResultType type, String data) {
+ this.type = type;
+ this.data = data;
+ }
+
+ public ResultType getType() {
+ return type;
+ }
+
+ public String getData() {
+ return data;
+ }
+ }
+
+ final class ContractResult {
+ private final ResultCode code;
+ private final List messages;
+ private final Map config;
+ private final Map guiParameters;
+ private final Map noteGuiParameters;
+
+ public ContractResult(
+ ResultCode code,
+ List messages,
+ Map config,
+ Map guiParameters,
+ Map noteGuiParameters) {
+ this.code = code;
+ this.messages = Collections.unmodifiableList(new ArrayList<>(messages));
+ this.config = immutableCopy(config);
+ this.guiParameters = immutableCopy(guiParameters);
+ this.noteGuiParameters = immutableCopy(noteGuiParameters);
+ }
+
+ public ResultCode getCode() {
+ return code;
+ }
+
+ public List getMessages() {
+ return messages;
+ }
+
+ public Map getConfig() {
+ return config;
+ }
+
+ public Map getGuiParameters() {
+ return guiParameters;
+ }
+
+ public Map getNoteGuiParameters() {
+ return noteGuiParameters;
+ }
+ }
+
+ final class ProbeSnapshot {
+ private final int constructorCalls;
+ private final int openCalls;
+ private final int closeCalls;
+ private final String completionBuffer;
+ private final int completionCursor;
+ private final String completionParagraphId;
+ private final String cancelNoteId;
+ private final String cancelParagraphId;
+ private final ContractContext lastInterpretContext;
+
+ public ProbeSnapshot(
+ int constructorCalls,
+ int openCalls,
+ int closeCalls,
+ String completionBuffer,
+ int completionCursor,
+ String completionParagraphId,
+ String cancelNoteId,
+ String cancelParagraphId,
+ ContractContext lastInterpretContext) {
+ this.constructorCalls = constructorCalls;
+ this.openCalls = openCalls;
+ this.closeCalls = closeCalls;
+ this.completionBuffer = completionBuffer;
+ this.completionCursor = completionCursor;
+ this.completionParagraphId = completionParagraphId;
+ this.cancelNoteId = cancelNoteId;
+ this.cancelParagraphId = cancelParagraphId;
+ this.lastInterpretContext = lastInterpretContext;
+ }
+
+ public int getConstructorCalls() {
+ return constructorCalls;
+ }
+
+ public int getOpenCalls() {
+ return openCalls;
+ }
+
+ public int getCloseCalls() {
+ return closeCalls;
+ }
+
+ public String getCompletionBuffer() {
+ return completionBuffer;
+ }
+
+ public int getCompletionCursor() {
+ return completionCursor;
+ }
+
+ public String getCompletionParagraphId() {
+ return completionParagraphId;
+ }
+
+ public String getCancelNoteId() {
+ return cancelNoteId;
+ }
+
+ public String getCancelParagraphId() {
+ return cancelParagraphId;
+ }
+
+ public ContractContext getLastInterpretContext() {
+ return lastInterpretContext;
+ }
+ }
+
+ interface ContractProbe {
+ ProbeSnapshot snapshot();
+
+ boolean awaitInterpretStarted(Duration timeout) throws InterruptedException;
+
+ boolean awaitCancelObserved(Duration timeout) throws InterruptedException;
+
+ void releaseInterpretation();
+ }
+
+ interface ContractFaults {
+ void makeTransportUnavailable();
+
+ void abortPendingCalls();
+ }
+
+ final class ContractFailure extends Exception {
+ private final FailureCategory category;
+
+ public ContractFailure(FailureCategory category, String message, Throwable cause) {
+ super(message, cause);
+ this.category = category;
+ }
+
+ public FailureCategory getCategory() {
+ return category;
+ }
+ }
+
+ private static Map immutableCopy(Map source) {
+ return Collections.unmodifiableMap(new LinkedHashMap<>(source));
+ }
+}
diff --git a/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/InterpreterRpcContractFixture.java b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/InterpreterRpcContractFixture.java
new file mode 100644
index 00000000000..b5cd7c761bc
--- /dev/null
+++ b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/InterpreterRpcContractFixture.java
@@ -0,0 +1,254 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.zeppelin.interpreter.remote;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.List;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.zeppelin.interpreter.Interpreter;
+import org.apache.zeppelin.interpreter.InterpreterContext;
+import org.apache.zeppelin.interpreter.InterpreterException;
+import org.apache.zeppelin.interpreter.InterpreterResult;
+import org.apache.zeppelin.interpreter.remote.InterpreterRpcContractDriver.ContractContext;
+import org.apache.zeppelin.interpreter.remote.InterpreterRpcContractDriver.ContractProbe;
+import org.apache.zeppelin.interpreter.remote.InterpreterRpcContractDriver.ProbeSnapshot;
+import org.apache.zeppelin.interpreter.thrift.InterpreterCompletion;
+import org.apache.zeppelin.user.AuthenticationInfo;
+
+/**
+ * Reusable executable fixture for Server-to-Interpreter control RPC contract drivers.
+ *
+ * The generated completion type remains here only while it is part of the public
+ * {@link Interpreter} API. The fixture and its probe lifecycle are independent of any transport
+ * adapter, so later drivers can execute the same interpreter and shared scenarios.
+ */
+public final class InterpreterRpcContractFixture {
+
+ public static final String PROBE_ID_PROPERTY = "zeppelin.interpreter.contract.probe.id";
+
+ private static final Duration INTERPRET_TIMEOUT = Duration.ofSeconds(10);
+ private static final ConcurrentMap PROBES = new ConcurrentHashMap<>();
+
+ private InterpreterRpcContractFixture() {
+ }
+
+ public static Handle create(String probeId) {
+ ProbeState probe = new ProbeState();
+ ProbeState previous = PROBES.putIfAbsent(probeId, probe);
+ if (previous != null) {
+ throw new IllegalArgumentException("Duplicate contract probe " + probeId);
+ }
+ return new Handle(probeId, probe);
+ }
+
+ /** Owns one isolated fixture probe and removes it from the registry when closed. */
+ public static final class Handle implements ContractProbe, AutoCloseable {
+ private final String probeId;
+ private final ProbeState probe;
+
+ private Handle(String probeId, ProbeState probe) {
+ this.probeId = probeId;
+ this.probe = probe;
+ }
+
+ public String getProbeId() {
+ return probeId;
+ }
+
+ public String getInterpreterClassName() {
+ return ContractInterpreter.class.getName();
+ }
+
+ @Override
+ public ProbeSnapshot snapshot() {
+ return new ProbeSnapshot(
+ probe.constructorCalls.get(),
+ probe.openCalls.get(),
+ probe.closeCalls.get(),
+ probe.completionBuffer.get(),
+ probe.completionCursor.get(),
+ probe.completionParagraphId.get(),
+ probe.cancelNoteId.get(),
+ probe.cancelParagraphId.get(),
+ probe.lastInterpretContext.get());
+ }
+
+ @Override
+ public boolean awaitInterpretStarted(Duration timeout) throws InterruptedException {
+ return probe.interpretStarted.await(timeout.toMillis(), TimeUnit.MILLISECONDS);
+ }
+
+ @Override
+ public boolean awaitCancelObserved(Duration timeout) throws InterruptedException {
+ return probe.cancelObserved.await(timeout.toMillis(), TimeUnit.MILLISECONDS);
+ }
+
+ @Override
+ public void releaseInterpretation() {
+ probe.releaseInterpret.countDown();
+ }
+
+ @Override
+ public void close() {
+ releaseInterpretation();
+ PROBES.remove(probeId, probe);
+ }
+ }
+
+ /** Interpreter implementation exercised by each transport contract driver. */
+ public static class ContractInterpreter extends Interpreter {
+ private final ProbeState probe;
+
+ public ContractInterpreter(Properties properties) {
+ super(properties);
+ String probeId = properties.getProperty(PROBE_ID_PROPERTY);
+ probe = PROBES.get(probeId);
+ if (probe == null) {
+ throw new IllegalStateException("Unknown contract probe " + probeId);
+ }
+ probe.constructorCalls.incrementAndGet();
+ }
+
+ @Override
+ public void open() {
+ probe.openCalls.incrementAndGet();
+ }
+
+ @Override
+ public void close() {
+ probe.closeCalls.incrementAndGet();
+ }
+
+ @Override
+ public InterpreterResult interpret(String statement, InterpreterContext context)
+ throws InterpreterException {
+ probe.lastInterpretContext.set(toContractContext(context));
+ if (statement.startsWith("echo:")) {
+ return new InterpreterResult(
+ InterpreterResult.Code.SUCCESS, statement.substring("echo:".length()));
+ }
+ if ("inspect-context".equals(statement)) {
+ context.getConfig().put("interpreter-config", "updated");
+ context.getGui().getParams().put("interpreter-gui", "updated");
+ context.getNoteGui().getParams().put("interpreter-note-gui", "updated");
+ return new InterpreterResult(
+ InterpreterResult.Code.SUCCESS, InterpreterResult.Type.TABLE, "context-result");
+ }
+ if ("throw-application-error".equals(statement)) {
+ throw new InterpreterException("contract interpret failure");
+ }
+ if ("block-until-cancelled".equals(statement)) {
+ context.setProgress(42);
+ probe.interpretStarted.countDown();
+ try {
+ if (!probe.releaseInterpret.await(
+ INTERPRET_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)) {
+ return new InterpreterResult(InterpreterResult.Code.ERROR, "cancel timed out");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return new InterpreterResult(InterpreterResult.Code.ERROR, "interpret interrupted");
+ }
+ return new InterpreterResult(
+ probe.cancelled.get()
+ ? InterpreterResult.Code.SUCCESS
+ : InterpreterResult.Code.ERROR,
+ probe.cancelled.get() ? "cancelled" : "released without cancellation");
+ }
+ return new InterpreterResult(InterpreterResult.Code.ERROR, "unsupported statement");
+ }
+
+ @Override
+ public void cancel(InterpreterContext context) {
+ probe.cancelNoteId.set(context.getNoteId());
+ probe.cancelParagraphId.set(context.getParagraphId());
+ probe.cancelled.set(true);
+ probe.cancelObserved.countDown();
+ probe.releaseInterpret.countDown();
+ }
+
+ @Override
+ public FormType getFormType() {
+ return FormType.NATIVE;
+ }
+
+ @Override
+ public int getProgress(InterpreterContext context) {
+ return 7;
+ }
+
+ @Override
+ public List completion(
+ String buffer, int cursor, InterpreterContext context) throws InterpreterException {
+ if ("throw-completion-error".equals(buffer)) {
+ throw new InterpreterException("contract completion failure");
+ }
+ probe.completionBuffer.set(buffer);
+ probe.completionCursor.set(cursor);
+ probe.completionParagraphId.set(context.getParagraphId());
+ return Collections.singletonList(
+ new InterpreterCompletion("select", "select *", "keyword"));
+ }
+
+ private ContractContext toContractContext(InterpreterContext context) {
+ AuthenticationInfo authenticationInfo = context.getAuthenticationInfo();
+ Set roles = authenticationInfo.getRoles() == null
+ ? Collections.emptySet()
+ : authenticationInfo.getRoles();
+ return new ContractContext(
+ context.getNoteId(),
+ context.getNoteName(),
+ context.getParagraphId(),
+ context.getReplName(),
+ context.getParagraphTitle(),
+ context.getParagraphText(),
+ authenticationInfo.getUser(),
+ roles,
+ authenticationInfo.getTicket(),
+ context.getLocalProperties(),
+ context.getConfig(),
+ context.getGui().getParams(),
+ context.getNoteGui().getParams());
+ }
+ }
+
+ private static final class ProbeState {
+ private final AtomicInteger constructorCalls = new AtomicInteger();
+ private final AtomicInteger openCalls = new AtomicInteger();
+ private final AtomicInteger closeCalls = new AtomicInteger();
+ private final AtomicReference completionBuffer = new AtomicReference<>();
+ private final AtomicInteger completionCursor = new AtomicInteger();
+ private final AtomicReference completionParagraphId = new AtomicReference<>();
+ private final AtomicReference cancelNoteId = new AtomicReference<>();
+ private final AtomicReference cancelParagraphId = new AtomicReference<>();
+ private final AtomicReference lastInterpretContext = new AtomicReference<>();
+ private final AtomicBoolean cancelled = new AtomicBoolean();
+ private final CountDownLatch interpretStarted = new CountDownLatch(1);
+ private final CountDownLatch cancelObserved = new CountDownLatch(1);
+ private final CountDownLatch releaseInterpret = new CountDownLatch(1);
+ }
+}
diff --git a/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/ThriftInterpreterRpcContractDriver.java b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/ThriftInterpreterRpcContractDriver.java
new file mode 100644
index 00000000000..a3a7f138e53
--- /dev/null
+++ b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/ThriftInterpreterRpcContractDriver.java
@@ -0,0 +1,392 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.zeppelin.interpreter.remote;
+
+import static org.mockito.Mockito.mock;
+
+import com.google.gson.Gson;
+import com.google.gson.reflect.TypeToken;
+import java.lang.reflect.Type;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import org.apache.thrift.TApplicationException;
+import org.apache.thrift.TException;
+import org.apache.thrift.protocol.TBinaryProtocol;
+import org.apache.thrift.transport.TSocket;
+import org.apache.thrift.transport.TTransportException;
+import org.apache.zeppelin.display.GUI;
+import org.apache.zeppelin.interpreter.thrift.InterpreterCompletion;
+import org.apache.zeppelin.interpreter.thrift.InterpreterRPCException;
+import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterContext;
+import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterResult;
+import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterResultMessage;
+import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterService;
+import org.apache.zeppelin.user.AuthenticationInfo;
+
+/** Runs the neutral control RPC contract through generated Thrift clients and a real socket. */
+public final class ThriftInterpreterRpcContractDriver implements InterpreterRpcContractDriver {
+
+ private static final String INTERPRETER_GROUP_ID = "contract-group";
+ private static final String SESSION_ID = "contract-session";
+ private static final String USER_NAME = "contract-user";
+ private static final String LOCAL_REPOSITORY_PROPERTY = "zeppelin.interpreter.localRepo";
+ private static final String FORCE_SHUTDOWN_PROPERTY = "zeppelin.interpreter.forceShutdown";
+ private static final int SOCKET_TIMEOUT_MILLIS = 10_000;
+ private static final Duration SERVER_TIMEOUT = Duration.ofSeconds(10);
+ private static final Type MAP_TYPE = new TypeToken>() { }.getType();
+ private static final Gson GSON = new Gson();
+
+ private final InterpreterSpec interpreterSpec;
+ private final InterpreterRpcContractFixture.Handle fixture;
+ private final ContractFaults faults = new ContractFaults() {
+ @Override
+ public void makeTransportUnavailable() {
+ closeTransport(controlTransport);
+ }
+
+ @Override
+ public void abortPendingCalls() {
+ closeTransport(executionTransport);
+ }
+ };
+
+ private RemoteInterpreterServer server;
+ private TSocket executionTransport;
+ private TSocket controlTransport;
+ private RemoteInterpreterService.Client executionClient;
+ private RemoteInterpreterService.Client controlClient;
+ private String previousLocalRepository;
+ private String previousForceShutdown;
+ private String previousProbeId;
+
+ ThriftInterpreterRpcContractDriver(Path localRepository, String probeId) {
+ this.fixture = InterpreterRpcContractFixture.create(probeId);
+
+ InterpreterRef interpreter =
+ new InterpreterRef(
+ INTERPRETER_GROUP_ID, SESSION_ID, fixture.getInterpreterClassName());
+ Map properties = new LinkedHashMap<>();
+ properties.put(LOCAL_REPOSITORY_PROPERTY, localRepository.toString());
+ properties.put(FORCE_SHUTDOWN_PROPERTY, "false");
+ properties.put(InterpreterRpcContractFixture.PROBE_ID_PROPERTY, fixture.getProbeId());
+ this.interpreterSpec = new InterpreterSpec(interpreter, USER_NAME, properties);
+ }
+
+ @Override
+ public InterpreterSpec interpreterSpec() {
+ return interpreterSpec;
+ }
+
+ @Override
+ public void start() throws Exception {
+ previousLocalRepository = System.getProperty(LOCAL_REPOSITORY_PROPERTY);
+ previousForceShutdown = System.getProperty(FORCE_SHUTDOWN_PROPERTY);
+ previousProbeId = System.getProperty(InterpreterRpcContractFixture.PROBE_ID_PROPERTY);
+
+ server =
+ new RemoteInterpreterServer(
+ "localhost",
+ RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces(),
+ ":",
+ INTERPRETER_GROUP_ID,
+ true);
+ server.intpEventClient = mock(RemoteInterpreterEventClient.class);
+ server.start();
+ awaitServerRunning();
+
+ executionTransport = openTransport();
+ controlTransport = openTransport();
+ executionClient =
+ new RemoteInterpreterService.Client(new TBinaryProtocol(executionTransport));
+ controlClient = new RemoteInterpreterService.Client(new TBinaryProtocol(controlTransport));
+ controlClient.init(Collections.emptyMap());
+ }
+
+ @Override
+ public void createInterpreter(InterpreterSpec spec) throws ContractFailure {
+ invoke(() -> {
+ controlClient.createInterpreter(
+ spec.getInterpreter().getInterpreterGroupId(),
+ spec.getInterpreter().getSessionId(),
+ spec.getInterpreter().getClassName(),
+ spec.getProperties(),
+ spec.getUserName());
+ return null;
+ });
+ }
+
+ @Override
+ public FormType getFormType(InterpreterRef interpreter) throws ContractFailure {
+ return invoke(() -> FormType.valueOf(
+ controlClient.getFormType(interpreter.getSessionId(), interpreter.getClassName())));
+ }
+
+ @Override
+ public int getProgress(InterpreterRef interpreter, ContractContext context)
+ throws ContractFailure {
+ return invoke(() -> controlClient.getProgress(
+ interpreter.getSessionId(), interpreter.getClassName(), toRemoteContext(context)));
+ }
+
+ @Override
+ public List completion(
+ InterpreterRef interpreter, String buffer, int cursor, ContractContext context)
+ throws ContractFailure {
+ return invoke(() -> {
+ List completions = new ArrayList<>();
+ for (InterpreterCompletion completion : controlClient.completion(
+ interpreter.getSessionId(),
+ interpreter.getClassName(),
+ buffer,
+ cursor,
+ toRemoteContext(context))) {
+ completions.add(
+ new ContractCompletion(
+ completion.getName(), completion.getValue(), completion.getMeta()));
+ }
+ return completions;
+ });
+ }
+
+ @Override
+ public ContractResult interpret(
+ InterpreterRef interpreter, String statement, ContractContext context)
+ throws ContractFailure {
+ return invoke(() -> toContractResult(executionClient.interpret(
+ interpreter.getSessionId(),
+ interpreter.getClassName(),
+ statement,
+ toRemoteContext(context))));
+ }
+
+ @Override
+ public JobStatus getStatus(InterpreterRef interpreter, String jobId) throws ContractFailure {
+ return invoke(() -> JobStatus.valueOf(
+ controlClient.getStatus(interpreter.getSessionId(), jobId)));
+ }
+
+ @Override
+ public void cancel(InterpreterRef interpreter, ContractContext context)
+ throws ContractFailure {
+ invoke(() -> {
+ controlClient.cancel(
+ interpreter.getSessionId(), interpreter.getClassName(), toRemoteContext(context));
+ return null;
+ });
+ }
+
+ @Override
+ public void closeInterpreter(InterpreterRef interpreter) throws ContractFailure {
+ invoke(() -> {
+ controlClient.close(interpreter.getSessionId(), interpreter.getClassName());
+ return null;
+ });
+ }
+
+ @Override
+ public ContractProbe probe() {
+ return fixture;
+ }
+
+ @Override
+ public ContractFaults faults() {
+ return faults;
+ }
+
+ void closeControlTransportForWireTest() {
+ closeTransport(controlTransport);
+ }
+
+ String getRawStatusForWireTest() throws TException {
+ return controlClient.getStatus(SESSION_ID, "wire-smoke-job");
+ }
+
+ String getRawMissingInterpreterFormTypeForWireTest() throws TException {
+ return controlClient.getFormType("missing-session", fixture.getInterpreterClassName());
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ fixture.releaseInterpretation();
+ closeTransport(executionTransport);
+ closeTransport(controlTransport);
+
+ if (server != null) {
+ try {
+ server.close(SESSION_ID, fixture.getInterpreterClassName());
+ } finally {
+ if (server.isRunning()) {
+ server.shutdown();
+ }
+ awaitServerStopped();
+ server.join(SOCKET_TIMEOUT_MILLIS);
+ if (server.isAlive()) {
+ throw new IllegalStateException("RemoteInterpreterServer did not terminate");
+ }
+ }
+ }
+ } finally {
+ try {
+ if (server != null) {
+ shutdownResultCleaner();
+ }
+ } finally {
+ restoreSystemProperty(LOCAL_REPOSITORY_PROPERTY, previousLocalRepository);
+ restoreSystemProperty(FORCE_SHUTDOWN_PROPERTY, previousForceShutdown);
+ restoreSystemProperty(
+ InterpreterRpcContractFixture.PROBE_ID_PROPERTY, previousProbeId);
+ fixture.close();
+ }
+ }
+ }
+
+ private TSocket openTransport() throws TTransportException {
+ TSocket transport = new TSocket("localhost", server.getPort(), SOCKET_TIMEOUT_MILLIS);
+ transport.open();
+ return transport;
+ }
+
+ private RemoteInterpreterContext toRemoteContext(ContractContext context) {
+ AuthenticationInfo authenticationInfo =
+ new AuthenticationInfo(
+ context.getUserName(),
+ new LinkedHashSet<>(context.getUserRoles()),
+ context.getUserTicket());
+ GUI gui = new GUI();
+ gui.setParams(new LinkedHashMap<>(context.getGuiParameters()));
+ GUI noteGui = new GUI();
+ noteGui.setParams(new LinkedHashMap<>(context.getNoteGuiParameters()));
+
+ RemoteInterpreterContext remoteContext = new RemoteInterpreterContext();
+ remoteContext.setNoteId(context.getNoteId());
+ remoteContext.setNoteName(context.getNoteName());
+ remoteContext.setParagraphId(context.getParagraphId());
+ remoteContext.setReplName(context.getReplName());
+ remoteContext.setParagraphTitle(context.getParagraphTitle());
+ remoteContext.setParagraphText(context.getParagraphText());
+ remoteContext.setAuthenticationInfo(authenticationInfo.toJson());
+ remoteContext.setConfig(GSON.toJson(context.getConfig()));
+ remoteContext.setGui(gui.toJson());
+ remoteContext.setNoteGui(noteGui.toJson());
+ remoteContext.setLocalProperties(new HashMap<>(context.getLocalProperties()));
+ return remoteContext;
+ }
+
+ private ContractResult toContractResult(RemoteInterpreterResult remoteResult) {
+ List messages = new ArrayList<>();
+ for (RemoteInterpreterResultMessage message : remoteResult.getMsg()) {
+ messages.add(
+ new ContractResultMessage(
+ ResultType.valueOf(message.getType()), message.getData()));
+ }
+ Map config = GSON.fromJson(remoteResult.getConfig(), MAP_TYPE);
+ Map guiParameters = GUI.fromJson(remoteResult.getGui()).getParams();
+ Map noteGuiParameters = GUI.fromJson(remoteResult.getNoteGui()).getParams();
+ return new ContractResult(
+ ResultCode.valueOf(remoteResult.getCode()),
+ messages,
+ config,
+ guiParameters,
+ noteGuiParameters);
+ }
+
+ private T invoke(ThriftCall call) throws ContractFailure {
+ try {
+ return call.execute();
+ } catch (InterpreterRPCException e) {
+ String message = e.getErrorMessage();
+ throw new ContractFailure(categoryForDeclaredFailure(message), message, e);
+ } catch (TTransportException e) {
+ throw new ContractFailure(FailureCategory.TRANSPORT_UNAVAILABLE, e.getMessage(), e);
+ } catch (TApplicationException e) {
+ throw new ContractFailure(FailureCategory.OPERATION_FAILED, e.getMessage(), e);
+ } catch (TException e) {
+ throw new ContractFailure(FailureCategory.OPERATION_FAILED, e.getMessage(), e);
+ }
+ }
+
+ private FailureCategory categoryForDeclaredFailure(String message) {
+ String normalized = message == null ? "" : message.toLowerCase(Locale.ROOT);
+ if (normalized.contains("not created")
+ || normalized.contains("not initialized")
+ || normalized.contains("not found")
+ || normalized.contains("no interpreter")) {
+ return FailureCategory.INTERPRETER_NOT_FOUND;
+ }
+ return FailureCategory.OPERATION_FAILED;
+ }
+
+ private void awaitServerRunning() throws InterruptedException {
+ long deadline = System.nanoTime() + SERVER_TIMEOUT.toNanos();
+ while (!server.isRunning() && System.nanoTime() < deadline) {
+ Thread.sleep(20);
+ }
+ if (!server.isRunning()) {
+ throw new IllegalStateException("RemoteInterpreterServer did not start");
+ }
+ }
+
+ private void awaitServerStopped() throws InterruptedException {
+ long deadline = System.nanoTime() + SERVER_TIMEOUT.toNanos();
+ while (server.isRunning() && System.nanoTime() < deadline) {
+ Thread.sleep(20);
+ }
+ if (server.isRunning()) {
+ throw new IllegalStateException("RemoteInterpreterServer did not stop");
+ }
+ }
+
+ private void shutdownResultCleaner() throws Exception {
+ java.lang.reflect.Field resultCleaner =
+ RemoteInterpreterServer.class.getDeclaredField("resultCleanService");
+ resultCleaner.setAccessible(true);
+ ScheduledExecutorService executor = (ScheduledExecutorService) resultCleaner.get(server);
+ executor.shutdownNow();
+ executor.awaitTermination(SERVER_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
+ }
+
+ private void closeTransport(TSocket transport) {
+ if (transport != null) {
+ transport.close();
+ }
+ }
+
+ private void restoreSystemProperty(String property, String previousValue) {
+ if (previousValue == null) {
+ System.clearProperty(property);
+ } else {
+ System.setProperty(property, previousValue);
+ }
+ }
+
+ @FunctionalInterface
+ private interface ThriftCall {
+ T execute() throws TException;
+ }
+}
diff --git a/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/ThriftInterpreterRpcContractTest.java b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/ThriftInterpreterRpcContractTest.java
new file mode 100644
index 00000000000..4e220fb9b88
--- /dev/null
+++ b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/ThriftInterpreterRpcContractTest.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.zeppelin.interpreter.remote;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Path;
+import org.apache.thrift.transport.TTransportException;
+import org.apache.zeppelin.interpreter.thrift.InterpreterRPCException;
+import org.junit.jupiter.api.Test;
+
+/** Runs the shared Server-to-Interpreter control RPC contract against the current Thrift wire. */
+public class ThriftInterpreterRpcContractTest extends AbstractInterpreterRpcContractTest {
+
+ private ThriftInterpreterRpcContractDriver thriftDriver;
+
+ @Override
+ protected InterpreterRpcContractDriver createDriver(Path localRepository, String probeId) {
+ thriftDriver = new ThriftInterpreterRpcContractDriver(localRepository, probeId);
+ return thriftDriver;
+ }
+
+ @Test
+ void shouldSurfaceClosedSocketAsRawThriftTransportFailure() {
+ thriftDriver.closeControlTransportForWireTest();
+ assertThrows(TTransportException.class, thriftDriver::getRawStatusForWireTest);
+ }
+
+ @Test
+ void shouldSurfaceDeclaredThriftInterpreterFailureOnTheWire() {
+ InterpreterRPCException failure = assertThrows(
+ InterpreterRPCException.class,
+ thriftDriver::getRawMissingInterpreterFormTypeForWireTest);
+ assertTrue(failure.getErrorMessage().contains("not initialized"));
+ }
+}