raft testing with akka

actorcreation.scala

system.actorOf(Props(new RaftActor with EventStreamAllMessages))

EventStreamAllMessages.scala

/**
 * Use for testing.
 * 
 * Forwards all messages received to the system's EventStream.
 * Use this to deterministically `awaitForLeaderElection` etc.
 */
trait EventStreamAllMessages {
  this: Actor =>

  override def aroundReceive(receive: Actor.Receive, msg: Any) = {
    context.system.eventStream.publish(msg.asInstanceOf[AnyRef])

    receive.applyOrElse(msg, unhandled)
  }
}

RaftAwait.scala

implicit val probe = TestProbe() // somewhere

def subscribeElectedLeader()(implicit probe: TestProbe): Unit =
  system.eventStream.subscribe(probe.ref, classOf[ElectedAsLeader])

def awaitElectedLeader()(implicit probe: TestProbe): Unit =
  probe.expectMsgClass(max = 3.seconds, classOf[ElectedAsLeader])

tests.scala

it should "elect initial Leader" in {
  // given
  info("Before election: ")
  infoMemberStates()

  // when
  awaitElectedLeader()
  info("After election: ")
  infoMemberStates()

  // then
  members.count(_.stateName == Leader) should equal (1)
  members.count(_.stateName == Candidate) should equal (0)
  members.count(_.stateName == Follower) should equal (4)
}

it should "elect replacement Leader if current Leader dies" in {
  // given
  infoMemberStates()

  // when
  val leaderToStop = leader.get
  leaderToStop.stop()
  info(s"Stopped leader: ${simpleName(leaderToStop)}")
  
  // then
  awaitElectedLeader()
  info("New leader elected: ")
  infoMemberStates()

  members.count(_.stateName == Leader) should equal (1)
  members.count(_.stateName == Candidate) should equal (0)
  members.count(_.stateName == Follower) should equal (3)
}