opendaylight/md-sal/sal-distributed-datastore/src/test/java/org/opendaylight/controller/cluster/datastore/ShardManagerTest.java

   1 /*
   2  * Copyright (c) 2014, 2015 Cisco Systems, Inc. and others.  All rights reserved.
   3  *
   4  * This program and the accompanying materials are made available under the
   5  * terms of the Eclipse Public License v1.0 which accompanies this distribution,
   6  * and is available at http://www.eclipse.org/legal/epl-v10.html
   7  */
   8
   9 package org.opendaylight.controller.cluster.datastore;
  10
  11 import static org.junit.Assert.assertEquals;
  12 import static org.junit.Assert.assertNull;
  13 import static org.junit.Assert.assertSame;
  14 import static org.junit.Assert.assertTrue;
  15 import static org.mockito.Mockito.mock;
  16 import static org.mockito.Mockito.never;
  17 import static org.mockito.Mockito.times;
  18 import static org.mockito.Mockito.verify;
  19 import akka.actor.ActorRef;
  20 import akka.actor.ActorSystem;
  21 import akka.actor.AddressFromURIString;
  22 import akka.actor.Props;
  23 import akka.cluster.Cluster;
  24 import akka.cluster.ClusterEvent;
  25 import akka.dispatch.Dispatchers;
  26 import akka.japi.Creator;
  27 import akka.pattern.Patterns;
  28 import akka.persistence.RecoveryCompleted;
  29 import akka.testkit.JavaTestKit;
  30 import akka.testkit.TestActorRef;
  31 import akka.util.Timeout;
  32 import com.google.common.base.Optional;
  33 import com.google.common.collect.ImmutableMap;
  34 import com.google.common.collect.ImmutableSet;
  35 import com.google.common.util.concurrent.Uninterruptibles;
  36 import com.typesafe.config.ConfigFactory;
  37 import java.util.Arrays;
  38 import java.util.List;
  39 import java.util.Map;
  40 import java.util.concurrent.CountDownLatch;
  41 import java.util.concurrent.TimeUnit;
  42 import org.junit.After;
  43 import org.junit.Before;
  44 import org.junit.Test;
  45 import org.mockito.Mock;
  46 import org.mockito.MockitoAnnotations;
  47 import org.opendaylight.controller.cluster.datastore.exceptions.NoShardLeaderException;
  48 import org.opendaylight.controller.cluster.datastore.exceptions.NotInitializedException;
  49 import org.opendaylight.controller.cluster.datastore.exceptions.PrimaryNotFoundException;
  50 import org.opendaylight.controller.cluster.datastore.identifiers.ShardIdentifier;
  51 import org.opendaylight.controller.cluster.datastore.identifiers.ShardManagerIdentifier;
  52 import org.opendaylight.controller.cluster.datastore.messages.ActorInitialized;
  53 import org.opendaylight.controller.cluster.datastore.messages.FindLocalShard;
  54 import org.opendaylight.controller.cluster.datastore.messages.FindPrimary;
  55 import org.opendaylight.controller.cluster.datastore.messages.LocalPrimaryShardFound;
  56 import org.opendaylight.controller.cluster.datastore.messages.LocalShardFound;
  57 import org.opendaylight.controller.cluster.datastore.messages.LocalShardNotFound;
  58 import org.opendaylight.controller.cluster.datastore.messages.PrimaryShardInfo;
  59 import org.opendaylight.controller.cluster.datastore.messages.RemotePrimaryShardFound;
  60 import org.opendaylight.controller.cluster.datastore.messages.ShardLeaderStateChanged;
  61 import org.opendaylight.controller.cluster.datastore.messages.SwitchShardBehavior;
  62 import org.opendaylight.controller.cluster.datastore.messages.UpdateSchemaContext;
  63 import org.opendaylight.controller.cluster.datastore.utils.MessageCollectorActor;
  64 import org.opendaylight.controller.cluster.datastore.utils.MockClusterWrapper;
  65 import org.opendaylight.controller.cluster.datastore.utils.MockConfiguration;
  66 import org.opendaylight.controller.cluster.datastore.utils.PrimaryShardInfoFutureCache;
  67 import org.opendaylight.controller.cluster.notifications.LeaderStateChanged;
  68 import org.opendaylight.controller.cluster.notifications.RegisterRoleChangeListener;
  69 import org.opendaylight.controller.cluster.notifications.RoleChangeNotification;
  70 import org.opendaylight.controller.cluster.raft.RaftState;
  71 import org.opendaylight.controller.cluster.raft.base.messages.FollowerInitialSyncUpStatus;
  72 import org.opendaylight.controller.cluster.raft.base.messages.SwitchBehavior;
  73 import org.opendaylight.controller.cluster.raft.utils.InMemoryJournal;
  74 import org.opendaylight.controller.md.cluster.datastore.model.TestModel;
  75 import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTree;
  76 import org.opendaylight.yangtools.yang.model.api.SchemaContext;
  77 import scala.concurrent.Await;
  78 import scala.concurrent.Future;
  79 import scala.concurrent.duration.FiniteDuration;
  80
  81 public class ShardManagerTest extends AbstractActorTest {
  82     private static int ID_COUNTER = 1;
  83
  84     private final String shardMrgIDSuffix = "config" + ID_COUNTER++;
  85     private final String shardMgrID = "shard-manager-" + shardMrgIDSuffix;
  86
  87     @Mock
  88     private static CountDownLatch ready;
  89
  90     private static TestActorRef<MessageCollectorActor> mockShardActor;
  91
  92     private static String mockShardName;
  93
  94     private final DatastoreContext.Builder datastoreContextBuilder = DatastoreContext.newBuilder().
  95             dataStoreType(shardMrgIDSuffix).shardInitializationTimeout(600, TimeUnit.MILLISECONDS)
  96                    .shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(6);
  97
  98     private static ActorRef newMockShardActor(ActorSystem system, String shardName, String memberName) {
  99         String name = new ShardIdentifier(shardName, memberName,"config").toString();
 100         return TestActorRef.create(system, Props.create(MessageCollectorActor.class), name);
 101     }
 102
 103     private final PrimaryShardInfoFutureCache primaryShardInfoCache = new PrimaryShardInfoFutureCache();
 104
 105     @Before
 106     public void setUp() {
 107         MockitoAnnotations.initMocks(this);
 108
 109         InMemoryJournal.clear();
 110
 111         if(mockShardActor == null) {
 112             mockShardName = new ShardIdentifier(Shard.DEFAULT_NAME, "member-1", "config").toString();
 113             mockShardActor = TestActorRef.create(getSystem(), Props.create(MessageCollectorActor.class), mockShardName);
 114         }
 115
 116         mockShardActor.underlyingActor().clear();
 117     }
 118
 119     @After
 120     public void tearDown() {
 121         InMemoryJournal.clear();
 122     }
 123
 124     private Props newShardMgrProps(boolean persistent) {
 125         return ShardManager.props(new MockClusterWrapper(), new MockConfiguration(),
 126                 datastoreContextBuilder.persistent(persistent).build(), ready, primaryShardInfoCache);
 127     }
 128
 129     private Props newPropsShardMgrWithMockShardActor() {
 130         return newPropsShardMgrWithMockShardActor("shardManager", mockShardActor, new MockClusterWrapper(),
 131                 new MockConfiguration());
 132     }
 133
 134     private Props newPropsShardMgrWithMockShardActor(final String name, final ActorRef shardActor,
 135             final ClusterWrapper clusterWrapper, final Configuration config) {
 136         Creator<ShardManager> creator = new Creator<ShardManager>() {
 137             private static final long serialVersionUID = 1L;
 138             @Override
 139             public ShardManager create() throws Exception {
 140                 return new ForwardingShardManager(clusterWrapper, config, datastoreContextBuilder.build(),
 141                         ready, name, shardActor, primaryShardInfoCache);
 142             }
 143         };
 144
 145         return Props.create(new DelegatingShardManagerCreator(creator)).withDispatcher(Dispatchers.DefaultDispatcherId());
 146     }
 147
 148     @Test
 149     public void testOnReceiveFindPrimaryForNonExistentShard() throws Exception {
 150         new JavaTestKit(getSystem()) {{
 151             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 152
 153             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 154
 155             shardManager.tell(new FindPrimary("non-existent", false), getRef());
 156
 157             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 158         }};
 159     }
 160
 161     @Test
 162     public void testOnReceiveFindPrimaryForLocalLeaderShard() throws Exception {
 163         new JavaTestKit(getSystem()) {{
 164             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 165
 166             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 167
 168             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 169             shardManager.tell(new ActorInitialized(), mockShardActor);
 170
 171             DataTree mockDataTree = mock(DataTree.class);
 172             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 173                     DataStoreVersions.CURRENT_VERSION), getRef());
 174
 175             MessageCollectorActor.expectFirstMatching(mockShardActor, RegisterRoleChangeListener.class);
 176             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
 177                     RaftState.Leader.name())), mockShardActor);
 178
 179             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 180
 181             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 182             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 183                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 184             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 185         }};
 186     }
 187
 188     @Test
 189     public void testOnReceiveFindPrimaryForNonLocalLeaderShardBeforeMemberUp() throws Exception {
 190         new JavaTestKit(getSystem()) {{
 191             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 192
 193             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 194             shardManager.tell(new ActorInitialized(), mockShardActor);
 195
 196             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 197             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 198             shardManager.tell(new RoleChangeNotification(memberId1,
 199                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 200             shardManager.tell(new LeaderStateChanged(memberId1, memberId2, DataStoreVersions.CURRENT_VERSION), mockShardActor);
 201
 202             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 203
 204             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 205         }};
 206     }
 207
 208     @Test
 209     public void testOnReceiveFindPrimaryForNonLocalLeaderShard() throws Exception {
 210         new JavaTestKit(getSystem()) {{
 211             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 212
 213             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 214             shardManager.tell(new ActorInitialized(), mockShardActor);
 215
 216             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 217             MockClusterWrapper.sendMemberUp(shardManager, "member-2", getRef().path().toString());
 218
 219             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 220             shardManager.tell(new RoleChangeNotification(memberId1,
 221                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 222             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
 223             shardManager.tell(new ShardLeaderStateChanged(memberId1, memberId2, Optional.<DataTree>absent(),
 224                     leaderVersion), mockShardActor);
 225
 226             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 227
 228             RemotePrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 229             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 230                     primaryFound.getPrimaryPath().contains("member-2-shard-default"));
 231             assertEquals("getPrimaryVersion", leaderVersion, primaryFound.getPrimaryVersion());
 232         }};
 233     }
 234
 235     @Test
 236     public void testOnReceiveFindPrimaryForUninitializedShard() throws Exception {
 237         new JavaTestKit(getSystem()) {{
 238             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 239
 240             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 241
 242             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 243         }};
 244     }
 245
 246     @Test
 247     public void testOnReceiveFindPrimaryForInitializedShardWithNoRole() throws Exception {
 248         new JavaTestKit(getSystem()) {{
 249             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 250
 251             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 252             shardManager.tell(new ActorInitialized(), mockShardActor);
 253
 254             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 255
 256             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 257         }};
 258     }
 259
 260     @Test
 261     public void testOnReceiveFindPrimaryForFollowerShardWithNoInitialLeaderId() throws Exception {
 262         new JavaTestKit(getSystem()) {{
 263             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 264
 265             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 266             shardManager.tell(new ActorInitialized(), mockShardActor);
 267
 268             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 269             shardManager.tell(new RoleChangeNotification(memberId,
 270                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 271
 272             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 273
 274             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 275
 276             DataTree mockDataTree = mock(DataTree.class);
 277             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 278                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
 279
 280             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 281
 282             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 283             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 284                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 285             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 286         }};
 287     }
 288
 289     @Test
 290     public void testOnReceiveFindPrimaryWaitForShardLeader() throws Exception {
 291         new JavaTestKit(getSystem()) {{
 292             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 293
 294             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 295
 296             // We're passing waitUntilInitialized = true to FindPrimary so the response should be
 297             // delayed until we send ActorInitialized and RoleChangeNotification.
 298             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 299
 300             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 301
 302             shardManager.tell(new ActorInitialized(), mockShardActor);
 303
 304             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 305
 306             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 307             shardManager.tell(new RoleChangeNotification(memberId,
 308                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor);
 309
 310             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 311
 312             DataTree mockDataTree = mock(DataTree.class);
 313             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 314                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
 315
 316             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 317             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 318                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 319             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 320
 321             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 322         }};
 323     }
 324
 325     @Test
 326     public void testOnReceiveFindPrimaryWaitForReadyWithUninitializedShard() throws Exception {
 327         new JavaTestKit(getSystem()) {{
 328             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 329
 330             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 331
 332             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 333
 334             expectMsgClass(duration("2 seconds"), NotInitializedException.class);
 335
 336             shardManager.tell(new ActorInitialized(), mockShardActor);
 337
 338             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 339         }};
 340     }
 341
 342     @Test
 343     public void testOnReceiveFindPrimaryWaitForReadyWithCandidateShard() throws Exception {
 344         new JavaTestKit(getSystem()) {{
 345             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 346
 347             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 348             shardManager.tell(new ActorInitialized(), mockShardActor);
 349             shardManager.tell(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
 350                     null, RaftState.Candidate.name()), mockShardActor);
 351
 352             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 353
 354             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 355         }};
 356     }
 357
 358     @Test
 359     public void testOnReceiveFindPrimaryWaitForReadyWithIsolatedLeaderShard() throws Exception {
 360         new JavaTestKit(getSystem()) {{
 361             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 362
 363             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 364             shardManager.tell(new ActorInitialized(), mockShardActor);
 365             shardManager.tell(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
 366                     null, RaftState.IsolatedLeader.name()), mockShardActor);
 367
 368             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 369
 370             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 371         }};
 372     }
 373
 374     @Test
 375     public void testOnReceiveFindPrimaryWaitForReadyWithNoRoleShard() throws Exception {
 376         new JavaTestKit(getSystem()) {{
 377             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 378
 379             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 380             shardManager.tell(new ActorInitialized(), mockShardActor);
 381
 382             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 383
 384             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 385         }};
 386     }
 387
 388     @Test
 389     public void testOnReceiveFindPrimaryForRemoteShard() throws Exception {
 390         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 391
 392         // Create an ActorSystem ShardManager actor for member-1.
 393
 394         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 395         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 396
 397         ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 398
 399         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 400                 newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 401                         new MockConfiguration()), shardManagerID);
 402
 403         // Create an ActorSystem ShardManager actor for member-2.
 404
 405         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 406
 407         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 408
 409         final ActorRef mockShardActor2 = newMockShardActor(system2, "astronauts", "member-2");
 410
 411         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 412                 put("default", Arrays.asList("member-1", "member-2")).
 413                 put("astronauts", Arrays.asList("member-2")).build());
 414
 415         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 416                 newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 417                         mockConfig2), shardManagerID);
 418
 419         new JavaTestKit(system1) {{
 420
 421             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 422             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 423
 424             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 425
 426             String memberId2 = "member-2-shard-astronauts-" + shardMrgIDSuffix;
 427             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
 428             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2,
 429                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardActor2);
 430             shardManager2.tell(new RoleChangeNotification(memberId2,
 431                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 432
 433             shardManager1.underlyingActor().waitForMemberUp();
 434
 435             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 436
 437             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 438             String path = found.getPrimaryPath();
 439             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-astronauts-config"));
 440             assertEquals("getPrimaryVersion", leaderVersion, found.getPrimaryVersion());
 441
 442             shardManager2.underlyingActor().verifyFindPrimary();
 443
 444             Cluster.get(system2).down(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 445
 446             shardManager1.underlyingActor().waitForMemberRemoved();
 447
 448             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 449
 450             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 451         }};
 452
 453         JavaTestKit.shutdownActorSystem(system1);
 454         JavaTestKit.shutdownActorSystem(system2);
 455     }
 456
 457     @Test
 458     public void testShardAvailabilityOnChangeOfMemberReachability() throws Exception {
 459         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 460
 461         // Create an ActorSystem ShardManager actor for member-1.
 462
 463         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 464         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 465
 466         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 467
 468         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 469             newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 470                 new MockConfiguration()), shardManagerID);
 471
 472         // Create an ActorSystem ShardManager actor for member-2.
 473
 474         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 475
 476         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 477
 478         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 479
 480         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 481             put("default", Arrays.asList("member-1", "member-2")).build());
 482
 483         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 484             newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 485                 mockConfig2), shardManagerID);
 486
 487         new JavaTestKit(system1) {{
 488
 489             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 490             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 491             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 492             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 493
 494             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 495             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 496             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 497                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 498             shardManager1.tell(new RoleChangeNotification(memberId1,
 499                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 500             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class)),
 501                     DataStoreVersions.CURRENT_VERSION),
 502                 mockShardActor2);
 503             shardManager2.tell(new RoleChangeNotification(memberId2,
 504                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 505             shardManager1.underlyingActor().waitForMemberUp();
 506
 507             shardManager1.tell(new FindPrimary("default", true), getRef());
 508
 509             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 510             String path = found.getPrimaryPath();
 511             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 512
 513             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 514                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 515
 516             shardManager1.underlyingActor().waitForUnreachableMember();
 517
 518             shardManager1.tell(new FindPrimary("default", true), getRef());
 519
 520             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 521
 522             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 523                 createReachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 524
 525             shardManager1.underlyingActor().waitForReachableMember();
 526
 527             shardManager1.tell(new FindPrimary("default", true), getRef());
 528
 529             RemotePrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 530             String path1 = found1.getPrimaryPath();
 531             assertTrue("Unexpected primary path " + path1, path1.contains("member-2-shard-default-config"));
 532
 533         }};
 534
 535         JavaTestKit.shutdownActorSystem(system1);
 536         JavaTestKit.shutdownActorSystem(system2);
 537     }
 538
 539     @Test
 540     public void testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange() throws Exception {
 541         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 542
 543         // Create an ActorSystem ShardManager actor for member-1.
 544
 545         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 546         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 547
 548         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 549
 550         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 551             newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 552                 new MockConfiguration()), shardManagerID);
 553
 554         // Create an ActorSystem ShardManager actor for member-2.
 555
 556         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 557
 558         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 559
 560         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 561
 562         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 563             put("default", Arrays.asList("member-1", "member-2")).build());
 564
 565         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 566             newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 567                 mockConfig2), shardManagerID);
 568
 569         new JavaTestKit(system1) {{
 570
 571             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 572             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 573             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 574             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 575
 576             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 577             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 578             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 579                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 580             shardManager1.tell(new RoleChangeNotification(memberId1,
 581                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 582             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class)),
 583                     DataStoreVersions.CURRENT_VERSION),
 584                 mockShardActor2);
 585             shardManager2.tell(new RoleChangeNotification(memberId2,
 586                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 587             shardManager1.underlyingActor().waitForMemberUp();
 588
 589             shardManager1.tell(new FindPrimary("default", true), getRef());
 590
 591             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 592             String path = found.getPrimaryPath();
 593             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 594
 595             primaryShardInfoCache.putSuccessful("default", new PrimaryShardInfo(system1.actorSelection(
 596                     mockShardActor1.path()), DataStoreVersions.CURRENT_VERSION, Optional.<DataTree>absent()));
 597
 598             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 599                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 600
 601             shardManager1.underlyingActor().waitForUnreachableMember();
 602
 603             shardManager1.tell(new FindPrimary("default", true), getRef());
 604
 605             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 606
 607             assertNull("Expected primaryShardInfoCache entry removed", primaryShardInfoCache.getIfPresent("default"));
 608
 609             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId1, Optional.of(mock(DataTree.class)),
 610                     DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 611             shardManager1.tell(new RoleChangeNotification(memberId1,
 612                 RaftState.Follower.name(), RaftState.Leader.name()), mockShardActor1);
 613
 614             shardManager1.tell(new FindPrimary("default", true), getRef());
 615
 616             LocalPrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 617             String path1 = found1.getPrimaryPath();
 618             assertTrue("Unexpected primary path " + path1, path1.contains("member-1-shard-default-config"));
 619
 620         }};
 621
 622         JavaTestKit.shutdownActorSystem(system1);
 623         JavaTestKit.shutdownActorSystem(system2);
 624     }
 625
 626
 627     @Test
 628     public void testOnReceiveFindLocalShardForNonExistentShard() throws Exception {
 629         new JavaTestKit(getSystem()) {{
 630             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 631
 632             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 633
 634             shardManager.tell(new FindLocalShard("non-existent", false), getRef());
 635
 636             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
 637
 638             assertEquals("getShardName", "non-existent", notFound.getShardName());
 639         }};
 640     }
 641
 642     @Test
 643     public void testOnReceiveFindLocalShardForExistentShard() throws Exception {
 644         new JavaTestKit(getSystem()) {{
 645             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 646
 647             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 648             shardManager.tell(new ActorInitialized(), mockShardActor);
 649
 650             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 651
 652             LocalShardFound found = expectMsgClass(duration("5 seconds"), LocalShardFound.class);
 653
 654             assertTrue("Found path contains " + found.getPath().path().toString(),
 655                     found.getPath().path().toString().contains("member-1-shard-default-config"));
 656         }};
 657     }
 658
 659     @Test
 660     public void testOnReceiveFindLocalShardForNotInitializedShard() throws Exception {
 661         new JavaTestKit(getSystem()) {{
 662             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 663
 664             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 665
 666             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 667         }};
 668     }
 669
 670     @Test
 671     public void testOnReceiveFindLocalShardWaitForShardInitialized() throws Exception {
 672         new JavaTestKit(getSystem()) {{
 673             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 674
 675             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 676
 677             // We're passing waitUntilInitialized = true to FindLocalShard so the response should be
 678             // delayed until we send ActorInitialized.
 679             Future<Object> future = Patterns.ask(shardManager, new FindLocalShard(Shard.DEFAULT_NAME, true),
 680                     new Timeout(5, TimeUnit.SECONDS));
 681
 682             shardManager.tell(new ActorInitialized(), mockShardActor);
 683
 684             Object resp = Await.result(future, duration("5 seconds"));
 685             assertTrue("Expected: LocalShardFound, Actual: " + resp, resp instanceof LocalShardFound);
 686         }};
 687     }
 688
 689     @Test
 690     public void testOnRecoveryJournalIsCleaned() {
 691         InMemoryJournal.addEntry(shardMgrID, 1L, new ShardManager.SchemaContextModules(
 692                 ImmutableSet.of("foo")));
 693         InMemoryJournal.addEntry(shardMgrID, 2L, new ShardManager.SchemaContextModules(
 694                 ImmutableSet.of("bar")));
 695         InMemoryJournal.addDeleteMessagesCompleteLatch(shardMgrID);
 696
 697         new JavaTestKit(getSystem()) {{
 698             TestActorRef<TestShardManager> shardManager = TestActorRef.create(getSystem(),
 699                     Props.create(new TestShardManagerCreator(shardMrgIDSuffix)));
 700
 701             shardManager.underlyingActor().waitForRecoveryComplete();
 702             InMemoryJournal.waitForDeleteMessagesComplete(shardMgrID);
 703
 704             // Journal entries up to the last one should've been deleted
 705             Map<Long, Object> journal = InMemoryJournal.get(shardMgrID);
 706             synchronized (journal) {
 707                 assertEquals("Journal size", 0, journal.size());
 708             }
 709         }};
 710     }
 711
 712     @Test
 713     public void testRoleChangeNotificationAndShardLeaderStateChangedReleaseReady() throws Exception {
 714         new JavaTestKit(getSystem()) {
 715             {
 716                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 717
 718                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 719                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 720                         memberId, RaftState.Candidate.name(), RaftState.Leader.name()));
 721
 722                 verify(ready, never()).countDown();
 723
 724                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId, memberId,
 725                         Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION));
 726
 727                 verify(ready, times(1)).countDown();
 728
 729             }};
 730     }
 731
 732     @Test
 733     public void testRoleChangeNotificationToFollowerWithShardLeaderStateChangedReleaseReady() throws Exception {
 734         new JavaTestKit(getSystem()) {
 735             {
 736                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 737
 738                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 739                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 740                         memberId, null, RaftState.Follower.name()));
 741
 742                 verify(ready, never()).countDown();
 743
 744                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 745
 746                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 747                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
 748                         DataStoreVersions.CURRENT_VERSION));
 749
 750                 verify(ready, times(1)).countDown();
 751
 752             }};
 753     }
 754
 755     @Test
 756     public void testReadyCountDownForMemberUpAfterLeaderStateChanged() throws Exception {
 757         new JavaTestKit(getSystem()) {
 758             {
 759                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 760
 761                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 762                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 763                         memberId, null, RaftState.Follower.name()));
 764
 765                 verify(ready, never()).countDown();
 766
 767                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 768                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
 769                         DataStoreVersions.CURRENT_VERSION));
 770
 771                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 772
 773                 verify(ready, times(1)).countDown();
 774
 775             }};
 776     }
 777
 778     @Test
 779     public void testRoleChangeNotificationDoNothingForUnknownShard() throws Exception {
 780         new JavaTestKit(getSystem()) {
 781             {
 782                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps(true));
 783
 784                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 785                         "unknown", RaftState.Candidate.name(), RaftState.Leader.name()));
 786
 787                 verify(ready, never()).countDown();
 788
 789             }};
 790     }
 791
 792
 793     @Test
 794     public void testByDefaultSyncStatusIsFalse() throws Exception{
 795         final Props persistentProps = newShardMgrProps(true);
 796         final TestActorRef<ShardManager> shardManager =
 797                 TestActorRef.create(getSystem(), persistentProps);
 798
 799         ShardManager shardManagerActor = shardManager.underlyingActor();
 800
 801         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 802     }
 803
 804     @Test
 805     public void testWhenShardIsLeaderSyncStatusIsTrue() throws Exception{
 806         final Props persistentProps = ShardManager.props(
 807                 new MockClusterWrapper(),
 808                 new MockConfiguration(),
 809                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 810         final TestActorRef<ShardManager> shardManager =
 811                 TestActorRef.create(getSystem(), persistentProps);
 812
 813         ShardManager shardManagerActor = shardManager.underlyingActor();
 814         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 815                 RaftState.Follower.name(), RaftState.Leader.name()));
 816
 817         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 818     }
 819
 820     @Test
 821     public void testWhenShardIsCandidateSyncStatusIsFalse() throws Exception{
 822         final Props persistentProps = newShardMgrProps(true);
 823         final TestActorRef<ShardManager> shardManager =
 824                 TestActorRef.create(getSystem(), persistentProps);
 825
 826         ShardManager shardManagerActor = shardManager.underlyingActor();
 827         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 828                 RaftState.Follower.name(), RaftState.Candidate.name()));
 829
 830         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 831
 832         // Send a FollowerInitialSyncStatus with status = true for the replica whose current state is candidate
 833         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 834
 835         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 836     }
 837
 838     @Test
 839     public void testWhenShardIsFollowerSyncStatusDependsOnFollowerInitialSyncStatus() throws Exception{
 840         final Props persistentProps = ShardManager.props(
 841                 new MockClusterWrapper(),
 842                 new MockConfiguration(),
 843                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 844         final TestActorRef<ShardManager> shardManager =
 845                 TestActorRef.create(getSystem(), persistentProps);
 846
 847         ShardManager shardManagerActor = shardManager.underlyingActor();
 848         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 849                 RaftState.Candidate.name(), RaftState.Follower.name()));
 850
 851         // Initially will be false
 852         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 853
 854         // Send status true will make sync status true
 855         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 856
 857         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 858
 859         // Send status false will make sync status false
 860         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(false, "member-1-shard-default-unknown"));
 861
 862         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 863
 864     }
 865
 866     @Test
 867     public void testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards() throws Exception{
 868         final Props persistentProps = ShardManager.props(
 869                 new MockClusterWrapper(),
 870                 new MockConfiguration() {
 871                     @Override
 872                     public List<String> getMemberShardNames(String memberName) {
 873                         return Arrays.asList("default", "astronauts");
 874                     }
 875                 },
 876                 DatastoreContext.newBuilder().persistent(true).build(), ready, primaryShardInfoCache);
 877         final TestActorRef<ShardManager> shardManager =
 878                 TestActorRef.create(getSystem(), persistentProps);
 879
 880         ShardManager shardManagerActor = shardManager.underlyingActor();
 881
 882         // Initially will be false
 883         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 884
 885         // Make default shard leader
 886         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 887                 RaftState.Follower.name(), RaftState.Leader.name()));
 888
 889         // default = Leader, astronauts is unknown so sync status remains false
 890         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 891
 892         // Make astronauts shard leader as well
 893         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
 894                 RaftState.Follower.name(), RaftState.Leader.name()));
 895
 896         // Now sync status should be true
 897         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 898
 899         // Make astronauts a Follower
 900         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
 901                 RaftState.Leader.name(), RaftState.Follower.name()));
 902
 903         // Sync status is not true
 904         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 905
 906         // Make the astronauts follower sync status true
 907         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-astronauts-unknown"));
 908
 909         // Sync status is now true
 910         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 911
 912     }
 913
 914     @Test
 915     public void testOnReceiveSwitchShardBehavior() throws Exception {
 916         new JavaTestKit(getSystem()) {{
 917             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 918
 919             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 920             shardManager.tell(new ActorInitialized(), mockShardActor);
 921
 922             shardManager.tell(new SwitchShardBehavior(mockShardName, "Leader", 1000), getRef());
 923
 924             SwitchBehavior switchBehavior = MessageCollectorActor.expectFirstMatching(mockShardActor, SwitchBehavior.class);
 925
 926             assertEquals(RaftState.Leader, switchBehavior.getNewState());
 927             assertEquals(1000, switchBehavior.getNewTerm());
 928         }};
 929     }
 930
 931
 932     private static class TestShardManager extends ShardManager {
 933         private final CountDownLatch recoveryComplete = new CountDownLatch(1);
 934
 935         TestShardManager(String shardMrgIDSuffix) {
 936             super(new MockClusterWrapper(), new MockConfiguration(),
 937                     DatastoreContext.newBuilder().dataStoreType(shardMrgIDSuffix).build(), ready,
 938                     new PrimaryShardInfoFutureCache());
 939         }
 940
 941         @Override
 942         public void handleRecover(Object message) throws Exception {
 943             try {
 944                 super.handleRecover(message);
 945             } finally {
 946                 if(message instanceof RecoveryCompleted) {
 947                     recoveryComplete.countDown();
 948                 }
 949             }
 950         }
 951
 952         void waitForRecoveryComplete() {
 953             assertEquals("Recovery complete", true,
 954                     Uninterruptibles.awaitUninterruptibly(recoveryComplete, 5, TimeUnit.SECONDS));
 955         }
 956     }
 957
 958     @SuppressWarnings("serial")
 959     static class TestShardManagerCreator implements Creator<TestShardManager> {
 960         String shardMrgIDSuffix;
 961
 962         TestShardManagerCreator(String shardMrgIDSuffix) {
 963             this.shardMrgIDSuffix = shardMrgIDSuffix;
 964         }
 965
 966         @Override
 967         public TestShardManager create() throws Exception {
 968             return new TestShardManager(shardMrgIDSuffix);
 969         }
 970
 971     }
 972
 973     private static class DelegatingShardManagerCreator implements Creator<ShardManager> {
 974         private static final long serialVersionUID = 1L;
 975         private final Creator<ShardManager> delegate;
 976
 977         public DelegatingShardManagerCreator(Creator<ShardManager> delegate) {
 978             this.delegate = delegate;
 979         }
 980
 981         @Override
 982         public ShardManager create() throws Exception {
 983             return delegate.create();
 984         }
 985     }
 986
 987     private static class ForwardingShardManager extends ShardManager {
 988         private CountDownLatch findPrimaryMessageReceived = new CountDownLatch(1);
 989         private CountDownLatch memberUpReceived = new CountDownLatch(1);
 990         private CountDownLatch memberRemovedReceived = new CountDownLatch(1);
 991         private CountDownLatch memberUnreachableReceived = new CountDownLatch(1);
 992         private CountDownLatch memberReachableReceived = new CountDownLatch(1);
 993         private final ActorRef shardActor;
 994         private final String name;
 995
 996         protected ForwardingShardManager(ClusterWrapper cluster, Configuration configuration,
 997                 DatastoreContext datastoreContext, CountDownLatch waitTillReadyCountdownLatch, String name,
 998                 ActorRef shardActor, PrimaryShardInfoFutureCache primaryShardInfoCache) {
 999             super(cluster, configuration, datastoreContext, waitTillReadyCountdownLatch, primaryShardInfoCache);
1000             this.shardActor = shardActor;
1001             this.name = name;
1002         }
1003
1004         @Override
1005         public void handleCommand(Object message) throws Exception {
1006             try{
1007                 super.handleCommand(message);
1008             } finally {
1009                 if(message instanceof FindPrimary) {
1010                     findPrimaryMessageReceived.countDown();
1011                 } else if(message instanceof ClusterEvent.MemberUp) {
1012                     String role = ((ClusterEvent.MemberUp)message).member().roles().head();
1013                     if(!getCluster().getCurrentMemberName().equals(role)) {
1014                         memberUpReceived.countDown();
1015                     }
1016                 } else if(message instanceof ClusterEvent.MemberRemoved) {
1017                     String role = ((ClusterEvent.MemberRemoved)message).member().roles().head();
1018                     if(!getCluster().getCurrentMemberName().equals(role)) {
1019                         memberRemovedReceived.countDown();
1020                     }
1021                 } else if(message instanceof ClusterEvent.UnreachableMember) {
1022                     String role = ((ClusterEvent.UnreachableMember)message).member().roles().head();
1023                     if(!getCluster().getCurrentMemberName().equals(role)) {
1024                         memberUnreachableReceived.countDown();
1025                     }
1026                 } else if(message instanceof ClusterEvent.ReachableMember) {
1027                     String role = ((ClusterEvent.ReachableMember)message).member().roles().head();
1028                     if(!getCluster().getCurrentMemberName().equals(role)) {
1029                         memberReachableReceived.countDown();
1030                     }
1031                 }
1032             }
1033         }
1034
1035         @Override
1036         public String persistenceId() {
1037             return name;
1038         }
1039
1040         @Override
1041         protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
1042             return shardActor;
1043         }
1044
1045         void waitForMemberUp() {
1046             assertEquals("MemberUp received", true,
1047                     Uninterruptibles.awaitUninterruptibly(memberUpReceived, 5, TimeUnit.SECONDS));
1048             memberUpReceived = new CountDownLatch(1);
1049         }
1050
1051         void waitForMemberRemoved() {
1052             assertEquals("MemberRemoved received", true,
1053                     Uninterruptibles.awaitUninterruptibly(memberRemovedReceived, 5, TimeUnit.SECONDS));
1054             memberRemovedReceived = new CountDownLatch(1);
1055         }
1056
1057         void waitForUnreachableMember() {
1058             assertEquals("UnreachableMember received", true,
1059                 Uninterruptibles.awaitUninterruptibly(memberUnreachableReceived, 5, TimeUnit.SECONDS
1060                 ));
1061             memberUnreachableReceived = new CountDownLatch(1);
1062         }
1063
1064         void waitForReachableMember() {
1065             assertEquals("ReachableMember received", true,
1066                 Uninterruptibles.awaitUninterruptibly(memberReachableReceived, 5, TimeUnit.SECONDS));
1067             memberReachableReceived = new CountDownLatch(1);
1068         }
1069
1070         void verifyFindPrimary() {
1071             assertEquals("FindPrimary received", true,
1072                     Uninterruptibles.awaitUninterruptibly(findPrimaryMessageReceived, 5, TimeUnit.SECONDS));
1073             findPrimaryMessageReceived = new CountDownLatch(1);
1074         }
1075     }
1076 }