Merge pull request #100 from lerna-stack/fix-crashing-replicated-entity-behavior-testkit

Fix TestKit throws "Shard received unexpected message" after the entity passivated

Author: tksugimoto

CHANGELOG.md

@@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   This change will show you deserialization warnings during the rolling update, it's safe to ignore. 
   For more details, see [Akka 2.6.16 release note](https://akka.io/blog/news/2021/08/19/akka-2.6.16-released#rolling-upgrades)
 
+### Fixed
+- TestKit throws "Shard received unexpected message" exception after the entity passivated [PR#100](https://github.com/lerna-stack/akka-entity-replication/pull/100)
+
 ## [v2.0.0] - 2021-07-16
 [v2.0.0]: https://github.com/lerna-stack/akka-entity-replication/compare/v1.0.0...v2.0.0
 

src/main/scala/lerna/akka/entityreplication/typed/internal/testkit/ReplicatedEntityBehaviorTestKitImpl.scala

@@ -3,6 +3,7 @@ package lerna.akka.entityreplication.typed.internal.testkit
 import akka.actor.testkit.typed.scaladsl.{ ActorTestKit, SerializationTestKit }
 import akka.actor.typed.{ ActorRef, Behavior }
 import akka.actor.typed.scaladsl.adapter._
+import lerna.akka.entityreplication.ReplicationRegion
 import lerna.akka.entityreplication.model.NormalizedEntityId
 import lerna.akka.entityreplication.raft.RaftProtocol
 import lerna.akka.entityreplication.raft.model.{ EntityEvent, NoOp, ReplicatedLog, Term }
@@ -11,6 +12,7 @@ import lerna.akka.entityreplication.typed.ClusterReplication.ShardCommand
 import lerna.akka.entityreplication.typed.{ ReplicatedEntityContext, ReplicatedEntityTypeKey }
 import lerna.akka.entityreplication.typed.testkit.ReplicatedEntityBehaviorTestKit
 
+import scala.annotation.tailrec
 import scala.util.control.NonFatal
 
 private[entityreplication] class ReplicatedEntityBehaviorTestKitImpl[Command, Event, State](
@@ -103,18 +105,22 @@ private[entityreplication] class ReplicatedEntityBehaviorTestKitImpl[Command, Ev
   }
 
   private[this] def handleReplicatedEvent(): Option[Event] = {
-    val shardCommand: Option[ShardCommand] =
-      try {
-        Option(shardProbe.receiveMessage())
-      } catch {
-        case NonFatal(_) => None // It wasn't replicated until the timeout
-      }
-    val replicateCommand =
+    @tailrec
+    def findReplicateCommand(): Option[RaftProtocol.Replicate] = {
+      val shardCommand: Option[ShardCommand] =
+        try {
+          Option(shardProbe.receiveMessage())
+        } catch {
+          case NonFatal(_) => None // It wasn't replicated until the timeout
+        }
       shardCommand match {
-        case Some(r: RaftProtocol.Replicate) => Option(r)
-        case Some(m)                         => throw new IllegalStateException(s"Shard received unexpected message [$m]")
-        case None                            => None
+        case Some(r: RaftProtocol.Replicate)      => Option(r)
+        case Some(r: ReplicationRegion.Passivate) => findReplicateCommand() // retry
+        case Some(m)                              => throw new IllegalStateException(s"Shard received unexpected message [$m]")
+        case None                                 => None
       }
+    }
+    val replicateCommand = findReplicateCommand()
     replicateCommand.foreach { replicate =>
       // side effects
       replicatedLog = replicatedLog.append(EntityEvent(Option(normalizedEntityId), replicate.event), term)

src/test/scala/lerna/akka/entityreplication/typed/testkit/ReplicatedEntityBehaviorTestKitSpec.scala

@@ -72,6 +72,17 @@ class ReplicatedEntityBehaviorTestKitSpec extends FlatSpec with Matchers with Be
 
   behavior of "ReplicatedEntityBehaviorTestKit"
 
+  it should "still work after passivation which is triggered by handling message with 'tell' pattern" in {
+    // check no exception raised
+    replicatedEntityTestKit.runCommand(TellAndPassivate())
+    replicatedEntityTestKit.runCommand(TellAndReplicateAEventAndPassivate(inc = 10))
+    replicatedEntityTestKit.runCommand(TellAndReplicateAEvent(inc = 10))
+  }
+  it should "still work after passivation which is triggered by handling message with 'ask' pattern" in {
+    // check no exception raised
+    replicatedEntityTestKit.runCommand(AskAndReplicateAEventAndPassivate(inc = 10, _))
+    replicatedEntityTestKit.runCommand(AskAndReplicateAEvent(inc = 10, _))
+  }
   it should "report the command that could not be serialized" in {
     val ex = intercept[AssertionError] {
       replicatedEntityTestKit.runCommand(FailCommandSerialization())
@@ -237,14 +248,19 @@ object ReplicatedEntityBehaviorTestKitSpec {
 
     sealed trait Command
     final case class AskAndReplicateAEvent(inc: Int, replyTo: ActorRef[Reply]) extends Command with KryoSerializable
-    final case class AskAndNoReply(replyTo: ActorRef[Reply])                   extends Command with KryoSerializable
-    final case class AskAndNoEvents(replyTo: ActorRef[Reply])                  extends Command with KryoSerializable
-    final case class TellAndNoEvents()                                         extends Command with KryoSerializable
-    final case class TellAndReplicateAEvent(inc: Int)                          extends Command with KryoSerializable
-    final case class FailCommandSerialization()                                extends Command
-    final case class FailReplySerialization(replyTo: ActorRef[Reply])          extends Command with KryoSerializable
-    final case class FailStateSerialization()                                  extends Command with KryoSerializable
-    final case class FailEventSerialization()                                  extends Command with KryoSerializable
+    final case class AskAndReplicateAEventAndPassivate(inc: Int, replyTo: ActorRef[Reply])
+        extends Command
+        with KryoSerializable
+    final case class AskAndNoReply(replyTo: ActorRef[Reply])          extends Command with KryoSerializable
+    final case class AskAndNoEvents(replyTo: ActorRef[Reply])         extends Command with KryoSerializable
+    final case class TellAndNoEvents()                                extends Command with KryoSerializable
+    final case class TellAndPassivate()                               extends Command with KryoSerializable
+    final case class TellAndReplicateAEvent(inc: Int)                 extends Command with KryoSerializable
+    final case class TellAndReplicateAEventAndPassivate(inc: Int)     extends Command with KryoSerializable
+    final case class FailCommandSerialization()                       extends Command
+    final case class FailReplySerialization(replyTo: ActorRef[Reply]) extends Command with KryoSerializable
+    final case class FailStateSerialization()                         extends Command with KryoSerializable
+    final case class FailEventSerialization()                         extends Command with KryoSerializable
 
     trait Reply
     final case class SerializableReply(inc: Int, total: Int) extends Reply with KryoSerializable
@@ -277,6 +293,15 @@ object ReplicatedEntityBehaviorTestKitSpec {
               }
               .thenReply(replyTo)(state => SerializableReply(inc, state.total))
 
+          case AskAndReplicateAEventAndPassivate(inc, replyTo) =>
+            Effect
+              .replicate(SerializableEvent(inc))
+              .thenRun { _: State =>
+                context.log.info("complete depositing")
+              }
+              .thenPassivate()
+              .thenReply(replyTo)(state => SerializableReply(inc, state.total))
+
           case AskAndNoReply(_) =>
             Effect.noReply
 
@@ -289,9 +314,15 @@ object ReplicatedEntityBehaviorTestKitSpec {
                 context.log.info("receive a command")
               }.thenNoReply()
 
+          case TellAndPassivate() =>
+            Effect.passivate().thenNoReply()
+
           case TellAndReplicateAEvent(inc: Int) =>
             Effect.replicate(SerializableEvent(inc)).thenNoReply()
 
+          case TellAndReplicateAEventAndPassivate(inc: Int) =>
+            Effect.replicate(SerializableEvent(inc)).thenPassivate().thenNoReply()
+
           case FailCommandSerialization() =>
             Effect.noReply