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.assertNotNull;
  13 import static org.junit.Assert.assertNull;
  14 import static org.junit.Assert.assertSame;
  15 import static org.junit.Assert.assertTrue;
  16 import static org.mockito.Mockito.mock;
  17 import static org.mockito.Mockito.never;
  18 import static org.mockito.Mockito.times;
  19 import static org.mockito.Mockito.verify;
  20 import akka.actor.ActorRef;
  21 import akka.actor.ActorSystem;
  22 import akka.actor.AddressFromURIString;
  23 import akka.actor.PoisonPill;
  24 import akka.actor.Props;
  25 import akka.actor.Status;
  26 import akka.cluster.Cluster;
  27 import akka.cluster.ClusterEvent;
  28 import akka.dispatch.Dispatchers;
  29 import akka.japi.Creator;
  30 import akka.pattern.Patterns;
  31 import akka.persistence.RecoveryCompleted;
  32 import akka.testkit.JavaTestKit;
  33 import akka.testkit.TestActorRef;
  34 import akka.util.Timeout;
  35 import com.google.common.base.Optional;
  36 import com.google.common.collect.ImmutableMap;
  37 import com.google.common.collect.ImmutableSet;
  38 import com.google.common.collect.Sets;
  39 import com.google.common.util.concurrent.Uninterruptibles;
  40 import com.typesafe.config.ConfigFactory;
  41 import java.net.URI;
  42 import java.util.AbstractMap;
  43 import java.util.Arrays;
  44 import java.util.Collection;
  45 import java.util.Collections;
  46 import java.util.HashMap;
  47 import java.util.List;
  48 import java.util.Map;
  49 import java.util.Map.Entry;
  50 import java.util.concurrent.CountDownLatch;
  51 import java.util.concurrent.TimeUnit;
  52 import org.junit.After;
  53 import org.junit.Before;
  54 import org.junit.Test;
  55 import org.mockito.Mock;
  56 import org.mockito.Mockito;
  57 import org.mockito.MockitoAnnotations;
  58 import org.opendaylight.controller.cluster.datastore.config.Configuration;
  59 import org.opendaylight.controller.cluster.datastore.config.ConfigurationImpl;
  60 import org.opendaylight.controller.cluster.datastore.config.EmptyModuleShardConfigProvider;
  61 import org.opendaylight.controller.cluster.datastore.config.ModuleShardConfiguration;
  62 import org.opendaylight.controller.cluster.datastore.exceptions.NoShardLeaderException;
  63 import org.opendaylight.controller.cluster.datastore.exceptions.NotInitializedException;
  64 import org.opendaylight.controller.cluster.datastore.exceptions.PrimaryNotFoundException;
  65 import org.opendaylight.controller.cluster.datastore.identifiers.ShardIdentifier;
  66 import org.opendaylight.controller.cluster.datastore.identifiers.ShardManagerIdentifier;
  67 import org.opendaylight.controller.cluster.datastore.messages.ActorInitialized;
  68 import org.opendaylight.controller.cluster.datastore.messages.AddShardReplica;
  69 import org.opendaylight.controller.cluster.datastore.messages.CreateShard;
  70 import org.opendaylight.controller.cluster.datastore.messages.CreateShardReply;
  71 import org.opendaylight.controller.cluster.datastore.messages.FindLocalShard;
  72 import org.opendaylight.controller.cluster.datastore.messages.FindPrimary;
  73 import org.opendaylight.controller.cluster.datastore.messages.LocalPrimaryShardFound;
  74 import org.opendaylight.controller.cluster.datastore.messages.LocalShardFound;
  75 import org.opendaylight.controller.cluster.datastore.messages.LocalShardNotFound;
  76 import org.opendaylight.controller.cluster.datastore.messages.PeerDown;
  77 import org.opendaylight.controller.cluster.datastore.messages.PeerUp;
  78 import org.opendaylight.controller.cluster.datastore.messages.PrimaryShardInfo;
  79 import org.opendaylight.controller.cluster.datastore.messages.RemotePrimaryShardFound;
  80 import org.opendaylight.controller.cluster.datastore.messages.RemoveShardReplica;
  81 import org.opendaylight.controller.cluster.datastore.messages.ShardLeaderStateChanged;
  82 import org.opendaylight.controller.cluster.datastore.messages.SwitchShardBehavior;
  83 import org.opendaylight.controller.cluster.datastore.messages.UpdateSchemaContext;
  84 import org.opendaylight.controller.cluster.datastore.utils.MockClusterWrapper;
  85 import org.opendaylight.controller.cluster.datastore.utils.MockConfiguration;
  86 import org.opendaylight.controller.cluster.datastore.utils.PrimaryShardInfoFutureCache;
  87 import org.opendaylight.controller.cluster.notifications.LeaderStateChanged;
  88 import org.opendaylight.controller.cluster.notifications.RegisterRoleChangeListener;
  89 import org.opendaylight.controller.cluster.notifications.RoleChangeNotification;
  90 import org.opendaylight.controller.cluster.raft.RaftState;
  91 import org.opendaylight.controller.cluster.raft.base.messages.FollowerInitialSyncUpStatus;
  92 import org.opendaylight.controller.cluster.raft.base.messages.SwitchBehavior;
  93 import org.opendaylight.controller.cluster.raft.messages.AddServer;
  94 import org.opendaylight.controller.cluster.raft.messages.AddServerReply;
  95 import org.opendaylight.controller.cluster.raft.messages.ServerChangeStatus;
  96 import org.opendaylight.controller.cluster.raft.utils.InMemoryJournal;
  97 import org.opendaylight.controller.cluster.raft.utils.MessageCollectorActor;
  98 import org.opendaylight.controller.md.cluster.datastore.model.TestModel;
  99 import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTree;
 100 import org.opendaylight.yangtools.yang.model.api.SchemaContext;
 101 import scala.concurrent.Await;
 102 import scala.concurrent.Future;
 103 import scala.concurrent.duration.FiniteDuration;
 104
 105 public class ShardManagerTest extends AbstractActorTest {
 106     private static int ID_COUNTER = 1;
 107
 108     private final String shardMrgIDSuffix = "config" + ID_COUNTER++;
 109     private final String shardMgrID = "shard-manager-" + shardMrgIDSuffix;
 110
 111     @Mock
 112     private static CountDownLatch ready;
 113
 114     private static TestActorRef<MessageCollectorActor> mockShardActor;
 115
 116     private static String mockShardName;
 117
 118     private final DatastoreContext.Builder datastoreContextBuilder = DatastoreContext.newBuilder().
 119             dataStoreType(shardMrgIDSuffix).shardInitializationTimeout(600, TimeUnit.MILLISECONDS)
 120                    .shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(6);
 121
 122     private static ActorRef newMockShardActor(ActorSystem system, String shardName, String memberName) {
 123         String name = new ShardIdentifier(shardName, memberName,"config").toString();
 124         return TestActorRef.create(system, Props.create(MessageCollectorActor.class), name);
 125     }
 126
 127     private final PrimaryShardInfoFutureCache primaryShardInfoCache = new PrimaryShardInfoFutureCache();
 128
 129     @Before
 130     public void setUp() {
 131         MockitoAnnotations.initMocks(this);
 132
 133         InMemoryJournal.clear();
 134
 135         if(mockShardActor == null) {
 136             mockShardName = new ShardIdentifier(Shard.DEFAULT_NAME, "member-1", "config").toString();
 137             mockShardActor = TestActorRef.create(getSystem(), Props.create(MessageCollectorActor.class), mockShardName);
 138         }
 139
 140         mockShardActor.underlyingActor().clear();
 141     }
 142
 143     @After
 144     public void tearDown() {
 145         InMemoryJournal.clear();
 146     }
 147
 148     private Props newShardMgrProps() {
 149         return newShardMgrProps(new MockConfiguration());
 150     }
 151
 152     private static DatastoreContextFactory newDatastoreContextFactory(DatastoreContext datastoreContext) {
 153         DatastoreContextFactory mockFactory = mock(DatastoreContextFactory.class);
 154         Mockito.doReturn(datastoreContext).when(mockFactory).getBaseDatastoreContext();
 155         Mockito.doReturn(datastoreContext).when(mockFactory).getShardDatastoreContext(Mockito.anyString());
 156         return mockFactory;
 157     }
 158
 159     private Props newShardMgrProps(Configuration config) {
 160         return ShardManager.props(new MockClusterWrapper(), config,
 161                 newDatastoreContextFactory(datastoreContextBuilder.build()), ready, primaryShardInfoCache);
 162     }
 163
 164     private Props newPropsShardMgrWithMockShardActor() {
 165         return newPropsShardMgrWithMockShardActor("shardManager", mockShardActor, new MockClusterWrapper(),
 166                 new MockConfiguration());
 167     }
 168
 169     private Props newPropsShardMgrWithMockShardActor(final String name, final ActorRef shardActor,
 170             final ClusterWrapper clusterWrapper, final Configuration config) {
 171         Creator<ShardManager> creator = new Creator<ShardManager>() {
 172             private static final long serialVersionUID = 1L;
 173             @Override
 174             public ShardManager create() throws Exception {
 175                 return new ForwardingShardManager(clusterWrapper, config, newDatastoreContextFactory(
 176                         datastoreContextBuilder.build()), ready, name, shardActor, primaryShardInfoCache);
 177             }
 178         };
 179
 180         return Props.create(new DelegatingShardManagerCreator(creator)).withDispatcher(Dispatchers.DefaultDispatcherId());
 181     }
 182
 183     @Test
 184     public void testPerShardDatastoreContext() throws Exception {
 185         final DatastoreContextFactory mockFactory = newDatastoreContextFactory(
 186                 datastoreContextBuilder.shardElectionTimeoutFactor(5).build());
 187
 188         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 189                 shardElectionTimeoutFactor(6).build()).when(mockFactory).getShardDatastoreContext("default");
 190
 191         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 192                 shardElectionTimeoutFactor(7).build()).when(mockFactory).getShardDatastoreContext("topology");
 193
 194         final MockConfiguration mockConfig = new MockConfiguration() {
 195             @Override
 196             public Collection<String> getMemberShardNames(String memberName) {
 197                 return Arrays.asList("default", "topology");
 198             }
 199
 200             @Override
 201             public Collection<String> getMembersFromShardName(String shardName) {
 202                 return Arrays.asList("member-1");
 203             }
 204         };
 205
 206         final TestActorRef<MessageCollectorActor> defaultShardActor = TestActorRef.create(getSystem(),
 207                 Props.create(MessageCollectorActor.class), "default");
 208         final TestActorRef<MessageCollectorActor> topologyShardActor = TestActorRef.create(getSystem(),
 209                 Props.create(MessageCollectorActor.class), "topology");
 210
 211         final Map<String, Entry<ActorRef, DatastoreContext>> shardInfoMap = Collections.synchronizedMap(
 212                 new HashMap<String, Entry<ActorRef, DatastoreContext>>());
 213         shardInfoMap.put("default", new AbstractMap.SimpleEntry<ActorRef, DatastoreContext>(defaultShardActor, null));
 214         shardInfoMap.put("topology", new AbstractMap.SimpleEntry<ActorRef, DatastoreContext>(topologyShardActor, null));
 215
 216         final CountDownLatch newShardActorLatch = new CountDownLatch(2);
 217         final Creator<ShardManager> creator = new Creator<ShardManager>() {
 218             private static final long serialVersionUID = 1L;
 219             @Override
 220             public ShardManager create() throws Exception {
 221                 return new ShardManager(new MockClusterWrapper(), mockConfig, mockFactory, ready, primaryShardInfoCache) {
 222                     @Override
 223                     protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
 224                         Entry<ActorRef, DatastoreContext> entry = shardInfoMap.get(info.getShardName());
 225                         ActorRef ref = null;
 226                         if(entry != null) {
 227                             ref = entry.getKey();
 228                             entry.setValue(info.getDatastoreContext());
 229                         }
 230
 231                         newShardActorLatch.countDown();
 232                         return ref;
 233                     }
 234                 };
 235             }
 236         };
 237
 238         JavaTestKit kit = new JavaTestKit(getSystem());
 239
 240         final ActorRef shardManager = getSystem().actorOf(Props.create(new DelegatingShardManagerCreator(creator)).
 241                     withDispatcher(Dispatchers.DefaultDispatcherId()));
 242
 243         shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), kit.getRef());
 244
 245         assertEquals("Shard actors created", true, newShardActorLatch.await(5, TimeUnit.SECONDS));
 246         assertEquals("getShardElectionTimeoutFactor", 6, shardInfoMap.get("default").getValue().
 247                 getShardElectionTimeoutFactor());
 248         assertEquals("getShardElectionTimeoutFactor", 7, shardInfoMap.get("topology").getValue().
 249                 getShardElectionTimeoutFactor());
 250
 251         DatastoreContextFactory newMockFactory = newDatastoreContextFactory(
 252                 datastoreContextBuilder.shardElectionTimeoutFactor(5).build());
 253         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 254                 shardElectionTimeoutFactor(66).build()).when(newMockFactory).getShardDatastoreContext("default");
 255
 256         Mockito.doReturn(DatastoreContext.newBuilderFrom(datastoreContextBuilder.build()).
 257                 shardElectionTimeoutFactor(77).build()).when(newMockFactory).getShardDatastoreContext("topology");
 258
 259         shardManager.tell(newMockFactory, kit.getRef());
 260
 261         DatastoreContext newContext = MessageCollectorActor.expectFirstMatching(defaultShardActor, DatastoreContext.class);
 262         assertEquals("getShardElectionTimeoutFactor", 66, newContext.getShardElectionTimeoutFactor());
 263
 264         newContext = MessageCollectorActor.expectFirstMatching(topologyShardActor, DatastoreContext.class);
 265         assertEquals("getShardElectionTimeoutFactor", 77, newContext.getShardElectionTimeoutFactor());
 266
 267         defaultShardActor.tell(PoisonPill.getInstance(), ActorRef.noSender());
 268         topologyShardActor.tell(PoisonPill.getInstance(), ActorRef.noSender());
 269     }
 270
 271     @Test
 272     public void testOnReceiveFindPrimaryForNonExistentShard() throws Exception {
 273         new JavaTestKit(getSystem()) {{
 274             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 275
 276             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 277
 278             shardManager.tell(new FindPrimary("non-existent", false), getRef());
 279
 280             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 281         }};
 282     }
 283
 284     @Test
 285     public void testOnReceiveFindPrimaryForLocalLeaderShard() throws Exception {
 286         new JavaTestKit(getSystem()) {{
 287             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 288
 289             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 290
 291             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 292             shardManager.tell(new ActorInitialized(), mockShardActor);
 293
 294             DataTree mockDataTree = mock(DataTree.class);
 295             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 296                     DataStoreVersions.CURRENT_VERSION), getRef());
 297
 298             MessageCollectorActor.expectFirstMatching(mockShardActor, RegisterRoleChangeListener.class);
 299             shardManager.tell((new RoleChangeNotification(memberId, RaftState.Candidate.name(),
 300                     RaftState.Leader.name())), mockShardActor);
 301
 302             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 303
 304             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 305             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 306                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 307             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 308         }};
 309     }
 310
 311     @Test
 312     public void testOnReceiveFindPrimaryForNonLocalLeaderShardBeforeMemberUp() throws Exception {
 313         new JavaTestKit(getSystem()) {{
 314             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 315
 316             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 317             shardManager.tell(new ActorInitialized(), mockShardActor);
 318
 319             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 320             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 321             shardManager.tell(new RoleChangeNotification(memberId1,
 322                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 323             shardManager.tell(new LeaderStateChanged(memberId1, memberId2, DataStoreVersions.CURRENT_VERSION), mockShardActor);
 324
 325             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 326
 327             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 328         }};
 329     }
 330
 331     @Test
 332     public void testOnReceiveFindPrimaryForNonLocalLeaderShard() throws Exception {
 333         new JavaTestKit(getSystem()) {{
 334             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 335
 336             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 337             shardManager.tell(new ActorInitialized(), mockShardActor);
 338
 339             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 340             MockClusterWrapper.sendMemberUp(shardManager, "member-2", getRef().path().toString());
 341
 342             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 343             shardManager.tell(new RoleChangeNotification(memberId1,
 344                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 345             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
 346             shardManager.tell(new ShardLeaderStateChanged(memberId1, memberId2, Optional.<DataTree>absent(),
 347                     leaderVersion), mockShardActor);
 348
 349             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 350
 351             RemotePrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 352             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 353                     primaryFound.getPrimaryPath().contains("member-2-shard-default"));
 354             assertEquals("getPrimaryVersion", leaderVersion, primaryFound.getPrimaryVersion());
 355         }};
 356     }
 357
 358     @Test
 359     public void testOnReceiveFindPrimaryForUninitializedShard() throws Exception {
 360         new JavaTestKit(getSystem()) {{
 361             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 362
 363             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 364
 365             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 366         }};
 367     }
 368
 369     @Test
 370     public void testOnReceiveFindPrimaryForInitializedShardWithNoRole() throws Exception {
 371         new JavaTestKit(getSystem()) {{
 372             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 373
 374             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 375             shardManager.tell(new ActorInitialized(), mockShardActor);
 376
 377             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 378
 379             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 380         }};
 381     }
 382
 383     @Test
 384     public void testOnReceiveFindPrimaryForFollowerShardWithNoInitialLeaderId() throws Exception {
 385         new JavaTestKit(getSystem()) {{
 386             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 387
 388             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 389             shardManager.tell(new ActorInitialized(), mockShardActor);
 390
 391             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 392             shardManager.tell(new RoleChangeNotification(memberId,
 393                     RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor);
 394
 395             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 396
 397             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 398
 399             DataTree mockDataTree = mock(DataTree.class);
 400             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 401                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
 402
 403             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, false), getRef());
 404
 405             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 406             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 407                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 408             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 409         }};
 410     }
 411
 412     @Test
 413     public void testOnReceiveFindPrimaryWaitForShardLeader() throws Exception {
 414         new JavaTestKit(getSystem()) {{
 415             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 416
 417             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 418
 419             // We're passing waitUntilInitialized = true to FindPrimary so the response should be
 420             // delayed until we send ActorInitialized and RoleChangeNotification.
 421             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 422
 423             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 424
 425             shardManager.tell(new ActorInitialized(), mockShardActor);
 426
 427             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 428
 429             String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 430             shardManager.tell(new RoleChangeNotification(memberId,
 431                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor);
 432
 433             expectNoMsg(FiniteDuration.create(150, TimeUnit.MILLISECONDS));
 434
 435             DataTree mockDataTree = mock(DataTree.class);
 436             shardManager.tell(new ShardLeaderStateChanged(memberId, memberId, Optional.of(mockDataTree),
 437                     DataStoreVersions.CURRENT_VERSION), mockShardActor);
 438
 439             LocalPrimaryShardFound primaryFound = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 440             assertTrue("Unexpected primary path " +  primaryFound.getPrimaryPath(),
 441                     primaryFound.getPrimaryPath().contains("member-1-shard-default"));
 442             assertSame("getLocalShardDataTree", mockDataTree, primaryFound.getLocalShardDataTree() );
 443
 444             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 445         }};
 446     }
 447
 448     @Test
 449     public void testOnReceiveFindPrimaryWaitForReadyWithUninitializedShard() throws Exception {
 450         new JavaTestKit(getSystem()) {{
 451             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 452
 453             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 454
 455             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 456
 457             expectMsgClass(duration("2 seconds"), NotInitializedException.class);
 458
 459             shardManager.tell(new ActorInitialized(), mockShardActor);
 460
 461             expectNoMsg(FiniteDuration.create(200, TimeUnit.MILLISECONDS));
 462         }};
 463     }
 464
 465     @Test
 466     public void testOnReceiveFindPrimaryWaitForReadyWithCandidateShard() throws Exception {
 467         new JavaTestKit(getSystem()) {{
 468             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 469
 470             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 471             shardManager.tell(new ActorInitialized(), mockShardActor);
 472             shardManager.tell(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
 473                     null, RaftState.Candidate.name()), mockShardActor);
 474
 475             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 476
 477             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 478         }};
 479     }
 480
 481     @Test
 482     public void testOnReceiveFindPrimaryWaitForReadyWithIsolatedLeaderShard() throws Exception {
 483         new JavaTestKit(getSystem()) {{
 484             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 485
 486             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 487             shardManager.tell(new ActorInitialized(), mockShardActor);
 488             shardManager.tell(new RoleChangeNotification("member-1-shard-default-" + shardMrgIDSuffix,
 489                     null, RaftState.IsolatedLeader.name()), mockShardActor);
 490
 491             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 492
 493             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 494         }};
 495     }
 496
 497     @Test
 498     public void testOnReceiveFindPrimaryWaitForReadyWithNoRoleShard() throws Exception {
 499         new JavaTestKit(getSystem()) {{
 500             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 501
 502             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 503             shardManager.tell(new ActorInitialized(), mockShardActor);
 504
 505             shardManager.tell(new FindPrimary(Shard.DEFAULT_NAME, true), getRef());
 506
 507             expectMsgClass(duration("2 seconds"), NoShardLeaderException.class);
 508         }};
 509     }
 510
 511     @Test
 512     public void testOnReceiveFindPrimaryForRemoteShard() throws Exception {
 513         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 514
 515         // Create an ActorSystem ShardManager actor for member-1.
 516
 517         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 518         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 519
 520         ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 521
 522         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 523                 newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 524                         new MockConfiguration()), shardManagerID);
 525
 526         // Create an ActorSystem ShardManager actor for member-2.
 527
 528         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 529
 530         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 531
 532         final ActorRef mockShardActor2 = newMockShardActor(system2, "astronauts", "member-2");
 533
 534         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 535                 put("default", Arrays.asList("member-1", "member-2")).
 536                 put("astronauts", Arrays.asList("member-2")).build());
 537
 538         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 539                 newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 540                         mockConfig2), shardManagerID);
 541
 542         new JavaTestKit(system1) {{
 543
 544             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 545             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 546
 547             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 548
 549             String memberId2 = "member-2-shard-astronauts-" + shardMrgIDSuffix;
 550             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
 551             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2,
 552                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardActor2);
 553             shardManager2.tell(new RoleChangeNotification(memberId2,
 554                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 555
 556             shardManager1.underlyingActor().waitForMemberUp();
 557
 558             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 559
 560             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 561             String path = found.getPrimaryPath();
 562             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-astronauts-config"));
 563             assertEquals("getPrimaryVersion", leaderVersion, found.getPrimaryVersion());
 564
 565             shardManager2.underlyingActor().verifyFindPrimary();
 566
 567             Cluster.get(system2).down(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 568
 569             shardManager1.underlyingActor().waitForMemberRemoved();
 570
 571             shardManager1.tell(new FindPrimary("astronauts", false), getRef());
 572
 573             expectMsgClass(duration("5 seconds"), PrimaryNotFoundException.class);
 574         }};
 575
 576         JavaTestKit.shutdownActorSystem(system1);
 577         JavaTestKit.shutdownActorSystem(system2);
 578     }
 579
 580     @Test
 581     public void testShardAvailabilityOnChangeOfMemberReachability() throws Exception {
 582         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 583
 584         // Create an ActorSystem ShardManager actor for member-1.
 585
 586         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 587         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 588
 589         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 590
 591         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 592             newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 593                 new MockConfiguration()), shardManagerID);
 594
 595         // Create an ActorSystem ShardManager actor for member-2.
 596
 597         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 598
 599         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 600
 601         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 602
 603         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 604             put("default", Arrays.asList("member-1", "member-2")).build());
 605
 606         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 607             newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 608                 mockConfig2), shardManagerID);
 609
 610         new JavaTestKit(system1) {{
 611
 612             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 613             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 614             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 615             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 616
 617             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 618             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 619             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 620                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 621             shardManager1.tell(new RoleChangeNotification(memberId1,
 622                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 623             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class)),
 624                     DataStoreVersions.CURRENT_VERSION),
 625                 mockShardActor2);
 626             shardManager2.tell(new RoleChangeNotification(memberId2,
 627                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 628             shardManager1.underlyingActor().waitForMemberUp();
 629
 630             shardManager1.tell(new FindPrimary("default", true), getRef());
 631
 632             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 633             String path = found.getPrimaryPath();
 634             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 635
 636             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 637                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 638
 639             shardManager1.underlyingActor().waitForUnreachableMember();
 640
 641             PeerDown peerDown = MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerDown.class);
 642             assertEquals("getMemberName", "member-2", peerDown.getMemberName());
 643             MessageCollectorActor.clearMessages(mockShardActor1);
 644
 645             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 646                     createMemberRemoved("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 647
 648             MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerDown.class);
 649
 650             shardManager1.tell(new FindPrimary("default", true), getRef());
 651
 652             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 653
 654             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 655                 createReachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 656
 657             shardManager1.underlyingActor().waitForReachableMember();
 658
 659             PeerUp peerUp = MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerUp.class);
 660             assertEquals("getMemberName", "member-2", peerUp.getMemberName());
 661             MessageCollectorActor.clearMessages(mockShardActor1);
 662
 663             shardManager1.tell(new FindPrimary("default", true), getRef());
 664
 665             RemotePrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 666             String path1 = found1.getPrimaryPath();
 667             assertTrue("Unexpected primary path " + path1, path1.contains("member-2-shard-default-config"));
 668
 669             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 670                     createMemberUp("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 671
 672             MessageCollectorActor.expectFirstMatching(mockShardActor1, PeerUp.class);
 673
 674         }};
 675
 676         JavaTestKit.shutdownActorSystem(system1);
 677         JavaTestKit.shutdownActorSystem(system2);
 678     }
 679
 680     @Test
 681     public void testShardAvailabilityChangeOnMemberUnreachableAndLeadershipChange() throws Exception {
 682         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
 683
 684         // Create an ActorSystem ShardManager actor for member-1.
 685
 686         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
 687         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 688
 689         final ActorRef mockShardActor1 = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
 690
 691         final TestActorRef<ForwardingShardManager> shardManager1 = TestActorRef.create(system1,
 692             newPropsShardMgrWithMockShardActor("shardManager1", mockShardActor1, new ClusterWrapperImpl(system1),
 693                 new MockConfiguration()), shardManagerID);
 694
 695         // Create an ActorSystem ShardManager actor for member-2.
 696
 697         final ActorSystem system2 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member2"));
 698
 699         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
 700
 701         final ActorRef mockShardActor2 = newMockShardActor(system2, Shard.DEFAULT_NAME, "member-2");
 702
 703         MockConfiguration mockConfig2 = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
 704             put("default", Arrays.asList("member-1", "member-2")).build());
 705
 706         final TestActorRef<ForwardingShardManager> shardManager2 = TestActorRef.create(system2,
 707             newPropsShardMgrWithMockShardActor("shardManager2", mockShardActor2, new ClusterWrapperImpl(system2),
 708                 mockConfig2), shardManagerID);
 709
 710         new JavaTestKit(system1) {{
 711
 712             shardManager1.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 713             shardManager2.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 714             shardManager1.tell(new ActorInitialized(), mockShardActor1);
 715             shardManager2.tell(new ActorInitialized(), mockShardActor2);
 716
 717             String memberId2 = "member-2-shard-default-" + shardMrgIDSuffix;
 718             String memberId1 = "member-1-shard-default-" + shardMrgIDSuffix;
 719             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId2,
 720                 Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 721             shardManager1.tell(new RoleChangeNotification(memberId1,
 722                 RaftState.Candidate.name(), RaftState.Follower.name()), mockShardActor1);
 723             shardManager2.tell(new ShardLeaderStateChanged(memberId2, memberId2, Optional.of(mock(DataTree.class)),
 724                     DataStoreVersions.CURRENT_VERSION),
 725                 mockShardActor2);
 726             shardManager2.tell(new RoleChangeNotification(memberId2,
 727                 RaftState.Candidate.name(), RaftState.Leader.name()), mockShardActor2);
 728             shardManager1.underlyingActor().waitForMemberUp();
 729
 730             shardManager1.tell(new FindPrimary("default", true), getRef());
 731
 732             RemotePrimaryShardFound found = expectMsgClass(duration("5 seconds"), RemotePrimaryShardFound.class);
 733             String path = found.getPrimaryPath();
 734             assertTrue("Unexpected primary path " + path, path.contains("member-2-shard-default-config"));
 735
 736             primaryShardInfoCache.putSuccessful("default", new PrimaryShardInfo(system1.actorSelection(
 737                     mockShardActor1.path()), DataStoreVersions.CURRENT_VERSION, Optional.<DataTree>absent()));
 738
 739             shardManager1.underlyingActor().onReceiveCommand(MockClusterWrapper.
 740                 createUnreachableMember("member-2", "akka.tcp://cluster-test@127.0.0.1:2558"));
 741
 742             shardManager1.underlyingActor().waitForUnreachableMember();
 743
 744             shardManager1.tell(new FindPrimary("default", true), getRef());
 745
 746             expectMsgClass(duration("5 seconds"), NoShardLeaderException.class);
 747
 748             assertNull("Expected primaryShardInfoCache entry removed", primaryShardInfoCache.getIfPresent("default"));
 749
 750             shardManager1.tell(new ShardLeaderStateChanged(memberId1, memberId1, Optional.of(mock(DataTree.class)),
 751                     DataStoreVersions.CURRENT_VERSION), mockShardActor1);
 752             shardManager1.tell(new RoleChangeNotification(memberId1,
 753                 RaftState.Follower.name(), RaftState.Leader.name()), mockShardActor1);
 754
 755             shardManager1.tell(new FindPrimary("default", true), getRef());
 756
 757             LocalPrimaryShardFound found1 = expectMsgClass(duration("5 seconds"), LocalPrimaryShardFound.class);
 758             String path1 = found1.getPrimaryPath();
 759             assertTrue("Unexpected primary path " + path1, path1.contains("member-1-shard-default-config"));
 760
 761         }};
 762
 763         JavaTestKit.shutdownActorSystem(system1);
 764         JavaTestKit.shutdownActorSystem(system2);
 765     }
 766
 767
 768     @Test
 769     public void testOnReceiveFindLocalShardForNonExistentShard() throws Exception {
 770         new JavaTestKit(getSystem()) {{
 771             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 772
 773             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 774
 775             shardManager.tell(new FindLocalShard("non-existent", false), getRef());
 776
 777             LocalShardNotFound notFound = expectMsgClass(duration("5 seconds"), LocalShardNotFound.class);
 778
 779             assertEquals("getShardName", "non-existent", notFound.getShardName());
 780         }};
 781     }
 782
 783     @Test
 784     public void testOnReceiveFindLocalShardForExistentShard() throws Exception {
 785         new JavaTestKit(getSystem()) {{
 786             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 787
 788             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 789             shardManager.tell(new ActorInitialized(), mockShardActor);
 790
 791             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 792
 793             LocalShardFound found = expectMsgClass(duration("5 seconds"), LocalShardFound.class);
 794
 795             assertTrue("Found path contains " + found.getPath().path().toString(),
 796                     found.getPath().path().toString().contains("member-1-shard-default-config"));
 797         }};
 798     }
 799
 800     @Test
 801     public void testOnReceiveFindLocalShardForNotInitializedShard() throws Exception {
 802         new JavaTestKit(getSystem()) {{
 803             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 804
 805             shardManager.tell(new FindLocalShard(Shard.DEFAULT_NAME, false), getRef());
 806
 807             expectMsgClass(duration("5 seconds"), NotInitializedException.class);
 808         }};
 809     }
 810
 811     @Test
 812     public void testOnReceiveFindLocalShardWaitForShardInitialized() throws Exception {
 813         new JavaTestKit(getSystem()) {{
 814             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
 815
 816             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
 817
 818             // We're passing waitUntilInitialized = true to FindLocalShard so the response should be
 819             // delayed until we send ActorInitialized.
 820             Future<Object> future = Patterns.ask(shardManager, new FindLocalShard(Shard.DEFAULT_NAME, true),
 821                     new Timeout(5, TimeUnit.SECONDS));
 822
 823             shardManager.tell(new ActorInitialized(), mockShardActor);
 824
 825             Object resp = Await.result(future, duration("5 seconds"));
 826             assertTrue("Expected: LocalShardFound, Actual: " + resp, resp instanceof LocalShardFound);
 827         }};
 828     }
 829
 830     @Test
 831     public void testOnRecoveryJournalIsCleaned() {
 832         InMemoryJournal.addEntry(shardMgrID, 1L, new ShardManager.SchemaContextModules(
 833                 ImmutableSet.of("foo")));
 834         InMemoryJournal.addEntry(shardMgrID, 2L, new ShardManager.SchemaContextModules(
 835                 ImmutableSet.of("bar")));
 836         InMemoryJournal.addDeleteMessagesCompleteLatch(shardMgrID);
 837
 838         new JavaTestKit(getSystem()) {{
 839             TestActorRef<TestShardManager> shardManager = TestActorRef.create(getSystem(),
 840                     Props.create(new TestShardManagerCreator(shardMrgIDSuffix)));
 841
 842             shardManager.underlyingActor().waitForRecoveryComplete();
 843             InMemoryJournal.waitForDeleteMessagesComplete(shardMgrID);
 844
 845             // Journal entries up to the last one should've been deleted
 846             Map<Long, Object> journal = InMemoryJournal.get(shardMgrID);
 847             synchronized (journal) {
 848                 assertEquals("Journal size", 0, journal.size());
 849             }
 850         }};
 851     }
 852
 853     @Test
 854     public void testRoleChangeNotificationAndShardLeaderStateChangedReleaseReady() throws Exception {
 855         new JavaTestKit(getSystem()) {
 856             {
 857                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps());
 858
 859                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 860                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 861                         memberId, RaftState.Candidate.name(), RaftState.Leader.name()));
 862
 863                 verify(ready, never()).countDown();
 864
 865                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId, memberId,
 866                         Optional.of(mock(DataTree.class)), DataStoreVersions.CURRENT_VERSION));
 867
 868                 verify(ready, times(1)).countDown();
 869
 870             }};
 871     }
 872
 873     @Test
 874     public void testRoleChangeNotificationToFollowerWithShardLeaderStateChangedReleaseReady() throws Exception {
 875         new JavaTestKit(getSystem()) {
 876             {
 877                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps());
 878
 879                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 880                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 881                         memberId, null, RaftState.Follower.name()));
 882
 883                 verify(ready, never()).countDown();
 884
 885                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 886
 887                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 888                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
 889                         DataStoreVersions.CURRENT_VERSION));
 890
 891                 verify(ready, times(1)).countDown();
 892
 893             }};
 894     }
 895
 896     @Test
 897     public void testReadyCountDownForMemberUpAfterLeaderStateChanged() throws Exception {
 898         new JavaTestKit(getSystem()) {
 899             {
 900                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps());
 901
 902                 String memberId = "member-1-shard-default-" + shardMrgIDSuffix;
 903                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 904                         memberId, null, RaftState.Follower.name()));
 905
 906                 verify(ready, never()).countDown();
 907
 908                 shardManager.underlyingActor().onReceiveCommand(new ShardLeaderStateChanged(memberId,
 909                         "member-2-shard-default-" + shardMrgIDSuffix, Optional.of(mock(DataTree.class)),
 910                         DataStoreVersions.CURRENT_VERSION));
 911
 912                 shardManager.underlyingActor().onReceiveCommand(MockClusterWrapper.createMemberUp("member-2", getRef().path().toString()));
 913
 914                 verify(ready, times(1)).countDown();
 915
 916             }};
 917     }
 918
 919     @Test
 920     public void testRoleChangeNotificationDoNothingForUnknownShard() throws Exception {
 921         new JavaTestKit(getSystem()) {
 922             {
 923                 TestActorRef<ShardManager> shardManager = TestActorRef.create(getSystem(), newShardMgrProps());
 924
 925                 shardManager.underlyingActor().onReceiveCommand(new RoleChangeNotification(
 926                         "unknown", RaftState.Candidate.name(), RaftState.Leader.name()));
 927
 928                 verify(ready, never()).countDown();
 929
 930             }};
 931     }
 932
 933
 934     @Test
 935     public void testByDefaultSyncStatusIsFalse() throws Exception{
 936         final Props persistentProps = newShardMgrProps();
 937         final TestActorRef<ShardManager> shardManager =
 938                 TestActorRef.create(getSystem(), persistentProps);
 939
 940         ShardManager shardManagerActor = shardManager.underlyingActor();
 941
 942         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 943     }
 944
 945     @Test
 946     public void testWhenShardIsLeaderSyncStatusIsTrue() throws Exception{
 947         final Props persistentProps = ShardManager.props(new MockClusterWrapper(), new MockConfiguration(),
 948                 newDatastoreContextFactory(DatastoreContext.newBuilder().persistent(true).build()), ready,
 949                 primaryShardInfoCache);
 950         final TestActorRef<ShardManager> shardManager =
 951                 TestActorRef.create(getSystem(), persistentProps);
 952
 953         ShardManager shardManagerActor = shardManager.underlyingActor();
 954         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 955                 RaftState.Follower.name(), RaftState.Leader.name()));
 956
 957         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 958     }
 959
 960     @Test
 961     public void testWhenShardIsCandidateSyncStatusIsFalse() throws Exception{
 962         final Props persistentProps = newShardMgrProps();
 963         final TestActorRef<ShardManager> shardManager =
 964                 TestActorRef.create(getSystem(), persistentProps);
 965
 966         ShardManager shardManagerActor = shardManager.underlyingActor();
 967         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 968                 RaftState.Follower.name(), RaftState.Candidate.name()));
 969
 970         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 971
 972         // Send a FollowerInitialSyncStatus with status = true for the replica whose current state is candidate
 973         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 974
 975         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 976     }
 977
 978     @Test
 979     public void testWhenShardIsFollowerSyncStatusDependsOnFollowerInitialSyncStatus() throws Exception{
 980         final Props persistentProps = ShardManager.props(new MockClusterWrapper(), new MockConfiguration(),
 981                 newDatastoreContextFactory(DatastoreContext.newBuilder().persistent(true).build()), ready,
 982                 primaryShardInfoCache);
 983         final TestActorRef<ShardManager> shardManager =
 984                 TestActorRef.create(getSystem(), persistentProps);
 985
 986         ShardManager shardManagerActor = shardManager.underlyingActor();
 987         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
 988                 RaftState.Candidate.name(), RaftState.Follower.name()));
 989
 990         // Initially will be false
 991         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
 992
 993         // Send status true will make sync status true
 994         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-default-unknown"));
 995
 996         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
 997
 998         // Send status false will make sync status false
 999         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(false, "member-1-shard-default-unknown"));
1000
1001         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
1002
1003     }
1004
1005     @Test
1006     public void testWhenMultipleShardsPresentSyncStatusMustBeTrueForAllShards() throws Exception{
1007         final Props persistentProps = ShardManager.props(new MockClusterWrapper(),
1008                 new MockConfiguration() {
1009                     @Override
1010                     public List<String> getMemberShardNames(String memberName) {
1011                         return Arrays.asList("default", "astronauts");
1012                     }
1013                 },
1014                 newDatastoreContextFactory(DatastoreContext.newBuilder().persistent(true).build()), ready,
1015                 primaryShardInfoCache);
1016         final TestActorRef<ShardManager> shardManager =
1017                 TestActorRef.create(getSystem(), persistentProps);
1018
1019         ShardManager shardManagerActor = shardManager.underlyingActor();
1020
1021         // Initially will be false
1022         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
1023
1024         // Make default shard leader
1025         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-default-unknown",
1026                 RaftState.Follower.name(), RaftState.Leader.name()));
1027
1028         // default = Leader, astronauts is unknown so sync status remains false
1029         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
1030
1031         // Make astronauts shard leader as well
1032         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
1033                 RaftState.Follower.name(), RaftState.Leader.name()));
1034
1035         // Now sync status should be true
1036         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
1037
1038         // Make astronauts a Follower
1039         shardManagerActor.onReceiveCommand(new RoleChangeNotification("member-1-shard-astronauts-unknown",
1040                 RaftState.Leader.name(), RaftState.Follower.name()));
1041
1042         // Sync status is not true
1043         assertEquals(false, shardManagerActor.getMBean().getSyncStatus());
1044
1045         // Make the astronauts follower sync status true
1046         shardManagerActor.onReceiveCommand(new FollowerInitialSyncUpStatus(true, "member-1-shard-astronauts-unknown"));
1047
1048         // Sync status is now true
1049         assertEquals(true, shardManagerActor.getMBean().getSyncStatus());
1050
1051     }
1052
1053     @Test
1054     public void testOnReceiveSwitchShardBehavior() throws Exception {
1055         new JavaTestKit(getSystem()) {{
1056             final ActorRef shardManager = getSystem().actorOf(newPropsShardMgrWithMockShardActor());
1057
1058             shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1059             shardManager.tell(new ActorInitialized(), mockShardActor);
1060
1061             shardManager.tell(new SwitchShardBehavior(mockShardName, "Leader", 1000), getRef());
1062
1063             SwitchBehavior switchBehavior = MessageCollectorActor.expectFirstMatching(mockShardActor, SwitchBehavior.class);
1064
1065             assertEquals(RaftState.Leader, switchBehavior.getNewState());
1066             assertEquals(1000, switchBehavior.getNewTerm());
1067         }};
1068     }
1069
1070     @Test
1071     public void testOnReceiveCreateShard() {
1072         new JavaTestKit(getSystem()) {{
1073             datastoreContextBuilder.shardInitializationTimeout(1, TimeUnit.MINUTES).persistent(true);
1074
1075             ActorRef shardManager = getSystem().actorOf(newShardMgrProps(
1076                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1077
1078             SchemaContext schemaContext = TestModel.createTestContext();
1079             shardManager.tell(new UpdateSchemaContext(schemaContext), ActorRef.noSender());
1080
1081             DatastoreContext datastoreContext = DatastoreContext.newBuilder().shardElectionTimeoutFactor(100).
1082                     persistent(false).build();
1083             TestShardPropsCreator shardPropsCreator = new TestShardPropsCreator();
1084
1085             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1086                     "foo", null, Arrays.asList("member-1", "member-5", "member-6"));
1087             shardManager.tell(new CreateShard(config, shardPropsCreator, datastoreContext), getRef());
1088
1089             expectMsgClass(duration("5 seconds"), CreateShardReply.class);
1090
1091             shardManager.tell(new FindLocalShard("foo", true), getRef());
1092
1093             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1094
1095             assertEquals("isRecoveryApplicable", false, shardPropsCreator.datastoreContext.isPersistent());
1096             assertTrue("Epxected ShardPeerAddressResolver", shardPropsCreator.datastoreContext.getShardRaftConfig().
1097                     getPeerAddressResolver() instanceof ShardPeerAddressResolver);
1098             assertEquals("peerMembers", Sets.newHashSet(new ShardIdentifier("foo", "member-5", shardMrgIDSuffix).toString(),
1099                     new ShardIdentifier("foo", "member-6", shardMrgIDSuffix).toString()),
1100                     shardPropsCreator.peerAddresses.keySet());
1101             assertEquals("ShardIdentifier", new ShardIdentifier("foo", "member-1", shardMrgIDSuffix),
1102                     shardPropsCreator.shardId);
1103             assertSame("schemaContext", schemaContext, shardPropsCreator.schemaContext);
1104
1105             // Send CreateShard with same name - should fail.
1106
1107             shardManager.tell(new CreateShard(config, shardPropsCreator, null), getRef());
1108
1109             expectMsgClass(duration("5 seconds"), akka.actor.Status.Failure.class);
1110         }};
1111     }
1112
1113     @Test
1114     public void testOnReceiveCreateShardWithNoInitialSchemaContext() {
1115         new JavaTestKit(getSystem()) {{
1116             ActorRef shardManager = getSystem().actorOf(newShardMgrProps(
1117                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1118
1119             TestShardPropsCreator shardPropsCreator = new TestShardPropsCreator();
1120
1121             ModuleShardConfiguration config = new ModuleShardConfiguration(URI.create("foo-ns"), "foo-module",
1122                     "foo", null, Arrays.asList("member-1"));
1123             shardManager.tell(new CreateShard(config, shardPropsCreator, null), getRef());
1124
1125             expectMsgClass(duration("5 seconds"), CreateShardReply.class);
1126
1127             SchemaContext schemaContext = TestModel.createTestContext();
1128             shardManager.tell(new UpdateSchemaContext(schemaContext), ActorRef.noSender());
1129
1130             shardManager.tell(new FindLocalShard("foo", true), getRef());
1131
1132             expectMsgClass(duration("5 seconds"), LocalShardFound.class);
1133
1134             assertSame("schemaContext", schemaContext, shardPropsCreator.schemaContext);
1135             assertNotNull("schemaContext is null", shardPropsCreator.datastoreContext);
1136         }};
1137     }
1138
1139     @Test
1140     public void testAddShardReplicaForNonExistentShard() throws Exception {
1141         new JavaTestKit(getSystem()) {{
1142             ActorRef shardManager = getSystem().actorOf(newShardMgrProps(
1143                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1144
1145             shardManager.tell(new AddShardReplica("model-inventory"), getRef());
1146             Status.Failure resp = expectMsgClass(duration("2 seconds"), Status.Failure.class);
1147
1148             assertEquals("Failure obtained", true,
1149                           (resp.cause() instanceof IllegalArgumentException));
1150         }};
1151     }
1152
1153     @Test
1154     public void testAddShardReplicaForAlreadyCreatedShard() throws Exception {
1155         new JavaTestKit(getSystem()) {{
1156             ActorRef shardManager = getSystem().actorOf(newShardMgrProps());
1157             shardManager.tell(new AddShardReplica("default"), getRef());
1158             Status.Failure resp = expectMsgClass(duration("2 seconds"), Status.Failure.class);
1159             assertEquals("Failure obtained", true,
1160                           (resp.cause() instanceof IllegalArgumentException));
1161         }};
1162     }
1163
1164     @Test
1165     public void testAddShardReplica() throws Exception {
1166         MockConfiguration mockConfig =
1167                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1168                    put("default", Arrays.asList("member-1", "member-2")).
1169                    put("astronauts", Arrays.asList("member-2")).build());
1170
1171         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
1172
1173         // Create an ActorSystem ShardManager actor for member-1.
1174         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
1175         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1176         ActorRef mockDefaultShardActor = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
1177         final TestActorRef<ForwardingShardManager> newReplicaShardManager = TestActorRef.create(system1,
1178                 newPropsShardMgrWithMockShardActor("shardManager1", mockDefaultShardActor,
1179                    new ClusterWrapperImpl(system1), mockConfig), shardManagerID);
1180
1181         // Create an ActorSystem ShardManager actor for member-2.
1182         final ActorSystem system2 = ActorSystem.create("cluster-test",
1183             ConfigFactory.load().getConfig("Member2"));
1184         Cluster.get(system2).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1185
1186         String name = new ShardIdentifier("astronauts", "member-2", "config").toString();
1187         final TestActorRef<MockRespondActor> mockShardLeaderActor =
1188             TestActorRef.create(system2, Props.create(MockRespondActor.class), name);
1189         final TestActorRef<ForwardingShardManager> leaderShardManager = TestActorRef.create(system2,
1190                 newPropsShardMgrWithMockShardActor("shardManager2", mockShardLeaderActor,
1191                         new ClusterWrapperImpl(system2), mockConfig), shardManagerID);
1192
1193         new JavaTestKit(system1) {{
1194
1195             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1196             leaderShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1197
1198             leaderShardManager.tell(new ActorInitialized(), mockShardLeaderActor);
1199
1200             String memberId2 = "member-2-shard-astronauts-" + shardMrgIDSuffix;
1201             short leaderVersion = DataStoreVersions.CURRENT_VERSION - 1;
1202             leaderShardManager.tell(new ShardLeaderStateChanged(memberId2, memberId2,
1203                     Optional.of(mock(DataTree.class)), leaderVersion), mockShardLeaderActor);
1204             leaderShardManager.tell(new RoleChangeNotification(memberId2,
1205                     RaftState.Candidate.name(), RaftState.Leader.name()), mockShardLeaderActor);
1206
1207             newReplicaShardManager.underlyingActor().waitForMemberUp();
1208             leaderShardManager.underlyingActor().waitForMemberUp();
1209
1210             //construct a mock response message
1211             AddServerReply response = new AddServerReply(ServerChangeStatus.OK, memberId2);
1212             mockShardLeaderActor.underlyingActor().updateResponse(response);
1213             newReplicaShardManager.tell(new AddShardReplica("astronauts"), getRef());
1214             AddServer addServerMsg = MessageCollectorActor.expectFirstMatching(mockShardLeaderActor,
1215                 AddServer.class);
1216             String addServerId = "member-1-shard-astronauts-" + shardMrgIDSuffix;
1217             assertEquals("AddServer serverId", addServerId, addServerMsg.getNewServerId());
1218
1219             expectMsgClass(duration("5 seconds"), Status.Success.class);
1220         }};
1221
1222         JavaTestKit.shutdownActorSystem(system1);
1223         JavaTestKit.shutdownActorSystem(system2);
1224     }
1225
1226     @Test
1227     public void testAddShardReplicaWithFindPrimaryTimeout() throws Exception {
1228         MockConfiguration mockConfig =
1229                 new MockConfiguration(ImmutableMap.<String, List<String>>builder().
1230                    put("default", Arrays.asList("member-1", "member-2")).
1231                    put("astronauts", Arrays.asList("member-2")).build());
1232
1233         String shardManagerID = ShardManagerIdentifier.builder().type(shardMrgIDSuffix).build().toString();
1234
1235         // Create an ActorSystem ShardManager actor for member-1.
1236         final ActorSystem system1 = ActorSystem.create("cluster-test", ConfigFactory.load().getConfig("Member1"));
1237         Cluster.get(system1).join(AddressFromURIString.parse("akka.tcp://cluster-test@127.0.0.1:2558"));
1238         ActorRef mockDefaultShardActor = newMockShardActor(system1, Shard.DEFAULT_NAME, "member-1");
1239         final TestActorRef<ForwardingShardManager> newReplicaShardManager = TestActorRef.create(system1,
1240                 newPropsShardMgrWithMockShardActor("shardManager1", mockDefaultShardActor,
1241                    new ClusterWrapperImpl(system1), mockConfig), shardManagerID);
1242
1243         new JavaTestKit(system1) {{
1244
1245             newReplicaShardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), getRef());
1246             MockClusterWrapper.sendMemberUp(newReplicaShardManager, "member-2", getRef().path().toString());
1247             newReplicaShardManager.underlyingActor().waitForMemberUp();
1248
1249             newReplicaShardManager.tell(new AddShardReplica("astronauts"), getRef());
1250             Status.Failure resp = expectMsgClass(duration("5 seconds"), Status.Failure.class);
1251             assertEquals("Failure obtained", true,
1252                           (resp.cause() instanceof RuntimeException));
1253         }};
1254
1255         JavaTestKit.shutdownActorSystem(system1);
1256     }
1257
1258     @Test
1259     public void testRemoveShardReplicaForNonExistentShard() throws Exception {
1260         new JavaTestKit(getSystem()) {{
1261             ActorRef shardManager = getSystem().actorOf(newShardMgrProps(
1262                     new ConfigurationImpl(new EmptyModuleShardConfigProvider())));
1263
1264             shardManager.tell(new RemoveShardReplica("model-inventory"), getRef());
1265             Status.Failure resp = expectMsgClass(duration("2 seconds"), Status.Failure.class);
1266             assertEquals("Failure obtained", true,
1267                          (resp.cause() instanceof IllegalArgumentException));
1268         }};
1269
1270     }
1271
1272     private static class TestShardPropsCreator implements ShardPropsCreator {
1273         ShardIdentifier shardId;
1274         Map<String, String> peerAddresses;
1275         SchemaContext schemaContext;
1276         DatastoreContext datastoreContext;
1277
1278         @Override
1279         public Props newProps(ShardIdentifier shardId, Map<String, String> peerAddresses,
1280                 DatastoreContext datastoreContext, SchemaContext schemaContext) {
1281             this.shardId = shardId;
1282             this.peerAddresses = peerAddresses;
1283             this.schemaContext = schemaContext;
1284             this.datastoreContext = datastoreContext;
1285             return Shard.props(shardId, peerAddresses, datastoreContext, schemaContext);
1286         }
1287
1288     }
1289
1290     private static class TestShardManager extends ShardManager {
1291         private final CountDownLatch recoveryComplete = new CountDownLatch(1);
1292
1293         TestShardManager(String shardMrgIDSuffix) {
1294             super(new MockClusterWrapper(), new MockConfiguration(),
1295                     newDatastoreContextFactory(DatastoreContext.newBuilder().dataStoreType(shardMrgIDSuffix).build()),
1296                     ready, new PrimaryShardInfoFutureCache());
1297         }
1298
1299         @Override
1300         public void handleRecover(Object message) throws Exception {
1301             try {
1302                 super.handleRecover(message);
1303             } finally {
1304                 if(message instanceof RecoveryCompleted) {
1305                     recoveryComplete.countDown();
1306                 }
1307             }
1308         }
1309
1310         void waitForRecoveryComplete() {
1311             assertEquals("Recovery complete", true,
1312                     Uninterruptibles.awaitUninterruptibly(recoveryComplete, 5, TimeUnit.SECONDS));
1313         }
1314     }
1315
1316     @SuppressWarnings("serial")
1317     static class TestShardManagerCreator implements Creator<TestShardManager> {
1318         String shardMrgIDSuffix;
1319
1320         TestShardManagerCreator(String shardMrgIDSuffix) {
1321             this.shardMrgIDSuffix = shardMrgIDSuffix;
1322         }
1323
1324         @Override
1325         public TestShardManager create() throws Exception {
1326             return new TestShardManager(shardMrgIDSuffix);
1327         }
1328
1329     }
1330
1331     private static class DelegatingShardManagerCreator implements Creator<ShardManager> {
1332         private static final long serialVersionUID = 1L;
1333         private final Creator<ShardManager> delegate;
1334
1335         public DelegatingShardManagerCreator(Creator<ShardManager> delegate) {
1336             this.delegate = delegate;
1337         }
1338
1339         @Override
1340         public ShardManager create() throws Exception {
1341             return delegate.create();
1342         }
1343     }
1344
1345     private static class ForwardingShardManager extends ShardManager {
1346         private CountDownLatch findPrimaryMessageReceived = new CountDownLatch(1);
1347         private CountDownLatch memberUpReceived = new CountDownLatch(1);
1348         private CountDownLatch memberRemovedReceived = new CountDownLatch(1);
1349         private CountDownLatch memberUnreachableReceived = new CountDownLatch(1);
1350         private CountDownLatch memberReachableReceived = new CountDownLatch(1);
1351         private final ActorRef shardActor;
1352         private final String name;
1353
1354         protected ForwardingShardManager(ClusterWrapper cluster, Configuration configuration,
1355                 DatastoreContextFactory factory, CountDownLatch waitTillReadyCountdownLatch, String name,
1356                 ActorRef shardActor, PrimaryShardInfoFutureCache primaryShardInfoCache) {
1357             super(cluster, configuration, factory, waitTillReadyCountdownLatch, primaryShardInfoCache);
1358             this.shardActor = shardActor;
1359             this.name = name;
1360         }
1361
1362         @Override
1363         public void handleCommand(Object message) throws Exception {
1364             try{
1365                 super.handleCommand(message);
1366             } finally {
1367                 if(message instanceof FindPrimary) {
1368                     findPrimaryMessageReceived.countDown();
1369                 } else if(message instanceof ClusterEvent.MemberUp) {
1370                     String role = ((ClusterEvent.MemberUp)message).member().roles().head();
1371                     if(!getCluster().getCurrentMemberName().equals(role)) {
1372                         memberUpReceived.countDown();
1373                     }
1374                 } else if(message instanceof ClusterEvent.MemberRemoved) {
1375                     String role = ((ClusterEvent.MemberRemoved)message).member().roles().head();
1376                     if(!getCluster().getCurrentMemberName().equals(role)) {
1377                         memberRemovedReceived.countDown();
1378                     }
1379                 } else if(message instanceof ClusterEvent.UnreachableMember) {
1380                     String role = ((ClusterEvent.UnreachableMember)message).member().roles().head();
1381                     if(!getCluster().getCurrentMemberName().equals(role)) {
1382                         memberUnreachableReceived.countDown();
1383                     }
1384                 } else if(message instanceof ClusterEvent.ReachableMember) {
1385                     String role = ((ClusterEvent.ReachableMember)message).member().roles().head();
1386                     if(!getCluster().getCurrentMemberName().equals(role)) {
1387                         memberReachableReceived.countDown();
1388                     }
1389                 }
1390             }
1391         }
1392
1393         @Override
1394         public String persistenceId() {
1395             return name;
1396         }
1397
1398         @Override
1399         protected ActorRef newShardActor(SchemaContext schemaContext, ShardInformation info) {
1400             return shardActor;
1401         }
1402
1403         void waitForMemberUp() {
1404             assertEquals("MemberUp received", true,
1405                     Uninterruptibles.awaitUninterruptibly(memberUpReceived, 5, TimeUnit.SECONDS));
1406             memberUpReceived = new CountDownLatch(1);
1407         }
1408
1409         void waitForMemberRemoved() {
1410             assertEquals("MemberRemoved received", true,
1411                     Uninterruptibles.awaitUninterruptibly(memberRemovedReceived, 5, TimeUnit.SECONDS));
1412             memberRemovedReceived = new CountDownLatch(1);
1413         }
1414
1415         void waitForUnreachableMember() {
1416             assertEquals("UnreachableMember received", true,
1417                 Uninterruptibles.awaitUninterruptibly(memberUnreachableReceived, 5, TimeUnit.SECONDS
1418                 ));
1419             memberUnreachableReceived = new CountDownLatch(1);
1420         }
1421
1422         void waitForReachableMember() {
1423             assertEquals("ReachableMember received", true,
1424                 Uninterruptibles.awaitUninterruptibly(memberReachableReceived, 5, TimeUnit.SECONDS));
1425             memberReachableReceived = new CountDownLatch(1);
1426         }
1427
1428         void verifyFindPrimary() {
1429             assertEquals("FindPrimary received", true,
1430                     Uninterruptibles.awaitUninterruptibly(findPrimaryMessageReceived, 5, TimeUnit.SECONDS));
1431             findPrimaryMessageReceived = new CountDownLatch(1);
1432         }
1433     }
1434
1435     private static class MockRespondActor extends MessageCollectorActor {
1436
1437         private volatile Object responseMsg;
1438
1439         public void updateResponse(Object response) {
1440             responseMsg = response;
1441         }
1442
1443         @Override
1444         public void onReceive(Object message) throws Exception {
1445             super.onReceive(message);
1446             if (message instanceof AddServer) {
1447                 if (responseMsg != null) {
1448                     getSender().tell(responseMsg, getSelf());
1449                     responseMsg = null;
1450                 }
1451             }
1452         }
1453     }
1454 }