From f07727d982fec26ceb25e8250e08b0188fd5d6c2 Mon Sep 17 00:00:00 2001
From: akarukappadath <39962314+akarukappadath@users.noreply.github.com>
Date: Wed, 7 Nov 2018 10:58:03 -0500
Subject: [PATCH] Retries for "Jupyter kernel is not ready" error (#668)

* trying retries

* cleaning up PR

* only retries if Jupyter kernel fails

* fixing compilation errors
---
 .../leonardo/LeonardoTestUtils.scala          | 19 +++++++++++++++----
 .../workbench/leonardo/NotebookPage.scala     | 11 ++++++++---
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/LeonardoTestUtils.scala b/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/LeonardoTestUtils.scala
index 7a17140c438..ae7a2843501 100644
--- a/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/LeonardoTestUtils.scala
+++ b/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/LeonardoTestUtils.scala
@@ -5,6 +5,7 @@ import java.nio.charset.StandardCharsets
 import java.nio.file.Files
 import java.util.Base64
 
+import akka.actor.ActorSystem
 import cats.data.OptionT
 import cats.implicits._
 import com.typesafe.scalalogging.LazyLogging
@@ -21,7 +22,7 @@ import org.broadinstitute.dsde.workbench.leonardo.Leonardo.ApiVersion.{V1, V2}
 import org.broadinstitute.dsde.workbench.leonardo.StringValueClass.LabelMap
 import org.broadinstitute.dsde.workbench.model.WorkbenchEmail
 import org.broadinstitute.dsde.workbench.model.google._
-import org.broadinstitute.dsde.workbench.util.LocalFileUtil
+import org.broadinstitute.dsde.workbench.util.{LocalFileUtil, Retry}
 import org.openqa.selenium.WebDriver
 import org.scalactic.source.Position
 import org.scalatest.{Matchers, Suite}
@@ -36,9 +37,10 @@ import scala.util.{Failure, Random, Success, Try}
 
 case class TimeResult[R](result:R, duration:FiniteDuration)
 
-trait LeonardoTestUtils extends WebBrowserSpec with Matchers with Eventually with LocalFileUtil with LazyLogging with ScalaFutures {
+trait LeonardoTestUtils extends WebBrowserSpec with Matchers with Eventually with LocalFileUtil with LazyLogging with ScalaFutures with Retry {
   this: Suite with BillingFixtures =>
 
+  val system: ActorSystem = ActorSystem("leotests")
   val logDir = new File("output")
   logDir.mkdirs
 
@@ -479,11 +481,20 @@ trait LeonardoTestUtils extends WebBrowserSpec with Matchers with Eventually wit
     }
   }
 
+  private def whenKernelNotReady(t: Throwable): Boolean = t match {
+    case e: KernelNotReadyException => true
+    case _ => false
+  }
+
+  implicit val ec: ExecutionContextExecutor = ExecutionContext.global
   def withNewNotebook[T](cluster: Cluster, kernel: Kernel = Python2, timeout: FiniteDuration = 2.minutes)(testCode: NotebookPage => T)(implicit webDriver: WebDriver, token: AuthToken): T = {
     withNotebooksListPage(cluster) { notebooksListPage =>
-      notebooksListPage.withNewNotebook(kernel, timeout) { notebookPage =>
-        testCode(notebookPage)
+      val result: Future[T] = retryUntilSuccessOrTimeout(whenKernelNotReady, failureLogMessage = s"Cannot make new notebook")(30 seconds, 2 minutes) {() =>
+        Future(notebooksListPage.withNewNotebook(kernel, timeout) { notebookPage =>
+          testCode(notebookPage)
+        })
       }
+      Await.result(result, 10 minutes)
     }
   }
 
diff --git a/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/NotebookPage.scala b/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/NotebookPage.scala
index a2a446a1b25..41602402233 100644
--- a/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/NotebookPage.scala
+++ b/automation/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/NotebookPage.scala
@@ -9,12 +9,15 @@ import org.scalatest.Assertions
 import org.scalatest.Matchers.convertToAnyShouldWrapper
 import org.scalatest.concurrent.Eventually
 import org.scalatest.concurrent.PatienceConfiguration.{Interval, Timeout}
+import org.scalatest.exceptions.TestFailedDueToTimeoutException
 import org.scalatest.time.{Seconds, Span}
 
 import scala.collection.JavaConverters._
 import scala.concurrent.duration._
 import scala.language.postfixOps
 
+case class KernelNotReadyException(timeElapsed:Timeout)
+  extends Exception(s"Jupyter kernel is NOT ready after waiting ${timeElapsed}")
 
 class NotebookPage(override val url: String)(override implicit val authToken: AuthToken, override implicit val webDriver: WebDriver)
   extends JupyterPage with Eventually with LazyLogging {
@@ -211,11 +214,13 @@ class NotebookPage(override val url: String)(override implicit val authToken: Au
   def awaitReadyKernel(timeout: FiniteDuration): Unit = {
     val time = Timeout(scaled(Span(timeout.toSeconds, Seconds)))
     val pollInterval = Interval(scaled(Span(5, Seconds)))
-    eventually(time, pollInterval) {
-      val ready = (!cellsAreRunning && isKernelReady && kernelNotificationText == "none")
-      Assertions.withClue(s"Jupyter kernel is NOT ready after waiting ${time}.") {
+    try {
+      eventually(time, pollInterval) {
+        val ready = (!cellsAreRunning && isKernelReady && kernelNotificationText == "none")
         ready shouldBe true
       }
+    } catch {
+      case e: TestFailedDueToTimeoutException => throw KernelNotReadyException(time)
     }
   }