HearbeatExpired is not stubbed correctly in test cases

Exception in thread "HeartbeatWatchdog-org.apache.livy.sessions.InteractiveSessionManager" org.mockito.exceptions.misusing.WrongTypeOfReturnValue: 
Idle$ cannot be returned by heartbeatExpired()
heartbeatExpired() should return boolean
***
If you're unsure why you're getting above error read on.
Due to the nature of the syntax above problem might occur because:
1. This exception *might* occur in wrongly written multi-threaded tests.
   Please refer to Mockito FAQ on limitations of concurrency testing.
2. A spy is stubbed using when(spy.foo()).then() syntax. It is safer to stub spies - 
   - with doReturn|Throw() family of methods. More in javadocs for Mockito.spy() method.

        at org.apache.livy.server.interactive.SessionHeartbeatWatchdog$$anonfun$deleteExpiredSessions$1.apply(SessionHeartbeat.scala:111)
        at org.apache.livy.server.interactive.SessionHeartbeatWatchdog$$anonfun$deleteExpiredSessions$1.apply(SessionHeartbeat.scala:110)
        at scala.collection.Iterator$class.foreach(Iterator.scala:891)
        at scala.collection.AbstractIterator.foreach(Iterator.scala:1334)
        at scala.collection.MapLike$DefaultValuesIterable.foreach(MapLike.scala:206)
        at org.apache.livy.server.interactive.SessionHeartbeatWatchdog$class.deleteExpiredSessions(SessionHeartbeat.scala:110)
        at org.apache.livy.sessions.InteractiveSessionManager.deleteExpiredSessions(SessionManager.scala:47)
        at org.apache.livy.server.interactive.SessionHeartbeatWatchdog$$anon$1.run(SessionHeartbeat.scala:92)
- interactive session should not gc-ed if session timeout check is off (1 second, 98 milliseconds)