add test for kernel die

python/src/test/java/org/apache/zeppelin/python/IPythonInterpreterTest.java

@@ -22,6 +22,7 @@ import org.apache.zeppelin.interpreter.InterpreterContext;
 import org.apache.zeppelin.interpreter.InterpreterException;
 import org.apache.zeppelin.interpreter.InterpreterGroup;
 import org.apache.zeppelin.interpreter.InterpreterResult;
+import org.apache.zeppelin.interpreter.InterpreterResult.Code;
 import org.apache.zeppelin.interpreter.InterpreterResultMessage;
 import org.apache.zeppelin.interpreter.LazyOpenInterpreter;
 import org.junit.Test;
@@ -65,6 +66,48 @@ public class IPythonInterpreterTest extends BasePythonInterpreterTest {
     intpGroup.close();
   }
 
+  @Test
+  public void testIpythonKernelCrash_shouldNotHangExecution()
+      throws InterpreterException, IOException {
+    // The goal of this test is to ensure that we handle case when the kernel die.
+    // In order to do so, we will kill the kernel process from the python code.
+    // A real example of that could be a out of memory by the code we execute.
+    String codeDep = "!pip install psutil";
+    String codeFindPID = "from os import getpid\n"
+        + "import psutil\n"
+        + "pids = psutil.pids()\n"
+        + "my_pid = getpid()\n"
+        + "pidToKill = []\n"
+        + "for pid in pids:\n"
+        + "    try:\n"
+        + "        p = psutil.Process(pid)\n"
+        + "        cmd = p.cmdline()\n"
+        + "        for arg in cmd:\n"
+        + "            if arg.count('ipykernel'):\n"
+        + "                pidToKill.append(pid)\n"
+        + "    except:\n"
+        + "        continue\n"
+        + "len(pidToKill)";
+    String codeKillKernel = "from os import kill\n"
+        + "import signal\n"
+        + "for pid in pidToKill:\n"
+        + "    kill(pid, signal.SIGKILL)";
+    InterpreterContext context = getInterpreterContext();
+    InterpreterResult result = interpreter.interpret(codeDep, context);
+    assertEquals(InterpreterResult.Code.SUCCESS, result.code());
+    context = getInterpreterContext();
+    result = interpreter.interpret(codeFindPID, context);
+    assertEquals(Code.SUCCESS, result.code());
+    InterpreterResultMessage output = context.out.toInterpreterResultMessage().get(0);
+    int numberOfPID = Integer.parseInt(output.getData());
+    assertTrue(numberOfPID > 0);
+    context = getInterpreterContext();
+    result = interpreter.interpret(codeKillKernel, context);
+    assertEquals(Code.ERROR, result.code());
+    output = context.out.toInterpreterResultMessage().get(0);
+    assertTrue(output.getData().equals("Ipython kernel has been stopped. Please check logs. It might be because of an out of memory issue."));
+  }
+
   @Test
   public void testIPythonAdvancedFeatures()
       throws InterpreterException, InterruptedException, IOException {