[LIVY-556] HearbeatExpired is not stubbed correctly in test cases

## What changes were proposed in this pull request? Add the proper stubbing to all mocked interactive sessions to make sure session heartbeats work correctly. ## How was this patch tested? This change only fixes the test cases. It is not changing production code. Author: Meisam Fathi <[email protected]> Closes #143 from meisam/LIVY-556.
apache · Feb 15, 2019 · 7d9b453 · 7d9b453
1 parent 2240c71
commit 7d9b453
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/server/src/test/scala/org/apache/livy/server/interactive/InteractiveSessionServletSpec.scala b/server/src/test/scala/org/apache/livy/server/interactive/InteractiveSessionServletSpec.scala
@@ -67,6 +67,7 @@ class InteractiveSessionServletSpec extends BaseInteractiveServletSpec {
       when(session.state).thenReturn(SessionState.Idle)
       when(session.stop()).thenReturn(Future.successful(()))
       when(session.proxyUser).thenReturn(None)
+      when(session.heartbeatExpired).thenReturn(false)
       when(session.statements).thenAnswer(
         new Answer[IndexedSeq[Statement]]() {
           override def answer(args: InvocationOnMock): IndexedSeq[Statement] = statements
@@ -179,6 +180,7 @@ class InteractiveSessionServletSpec extends BaseInteractiveServletSpec {
     when(session.kind).thenReturn(kind)
     when(session.appInfo).thenReturn(appInfo)
     when(session.logLines()).thenReturn(log)
+    when(session.heartbeatExpired).thenReturn(false)
 
     val req = mock[HttpServletRequest]
 

diff --git a/server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala b/server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala
@@ -99,6 +99,7 @@ class SessionManagerSpec extends FunSpec with Matchers with LivyBaseUnitTestSuit
       when(session.name).thenReturn(None)
       when(session.stop()).thenReturn(Future {})
       when(session.lastActivity).thenReturn(System.nanoTime())
+      when(session.heartbeatExpired).thenReturn(false)
 
       val conf = new LivyConf().set(LivyConf.SESSION_TIMEOUT_CHECK, false)
         .set(LivyConf.SESSION_STATE_RETAIN_TIME, "1s")