+ @Test
+ public void testLocalCandidateRemovedWithCandidateRegistered() throws Exception {
+ ShardTestKit kit = new ShardTestKit(getSystem());
+
+ dataStoreContextBuilder.shardHeartbeatIntervalInMillis(100).shardElectionTimeoutFactor(10000);
+ ShardIdentifier leaderId = newShardId("leader");
+ ShardIdentifier localId = newShardId(LOCAL_MEMBER_NAME);
+
+ TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(Props.create(
+ TestEntityOwnershipShard.class, localId,
+ ImmutableMap.<String, String>builder().put(leaderId.toString(), "".toString()).build(),
+ dataStoreContextBuilder.build()).withDispatcher(Dispatchers.DefaultDispatcherId()));
+
+ TestActorRef<EntityOwnershipShard> leader = actorFactory.createTestActor(newShardProps(leaderId,
+ ImmutableMap.<String, String>builder().put(localId.toString(), shard.path().toString()).build(),
+ LOCAL_MEMBER_NAME, EntityOwnerSelectionStrategyConfig.newBuilder().build()).withDispatcher(Dispatchers.DefaultDispatcherId()), leaderId.toString());
+ leader.tell(ElectionTimeout.INSTANCE, leader);
+
+ ShardTestKit.waitUntilLeader(leader);
+
+ shard.tell(new PeerAddressResolved(leaderId.toString(), leader.path().toString()), ActorRef.noSender());
+
+ Entity entity = new Entity(ENTITY_TYPE, ENTITY_ID1);
+ EntityOwnershipListener listener = mock(EntityOwnershipListener.class);
+
+ shard.tell(new RegisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ // Register local candidate
+
+ shard.tell(new RegisterCandidateLocal(entity), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+ verifyCommittedEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
+ verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity, false, true, true));
+ reset(listener);
+
+ // Simulate a replicated commit from the leader to remove the local candidate that would occur after a
+ // network partition is healed.
+
+ leader.tell(new PeerDown(localId.getMemberName(), localId.toString()), ActorRef.noSender());
+
+ verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity, true, false, false));
+
+ // Since the the shard has a local candidate registered, it should re-add its candidate to the entity.
+
+ verifyCommittedEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
+ verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity, false, true, true));
+
+ // Unregister the local candidate and verify it's removed and no re-added.
+
+ shard.tell(new UnregisterCandidateLocal(entity), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ verifyNoEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
+ Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
+ verifyNoEntityCandidate(shard, entity.getType(), entity.getId(), LOCAL_MEMBER_NAME);
+ }
+
+ @Test
+ public void testListenerRegistration() throws Exception {
+ ShardTestKit kit = new ShardTestKit(getSystem());
+ TestActorRef<EntityOwnershipShard> shard = actorFactory.createTestActor(newShardProps());
+ ShardTestKit.waitUntilLeader(shard);
+ ShardDataTree shardDataTree = shard.underlyingActor().getDataStore();
+
+ String otherEntityType = "otherEntityType";
+ Entity entity1 = new Entity(ENTITY_TYPE, ENTITY_ID1);
+ Entity entity2 = new Entity(ENTITY_TYPE, ENTITY_ID2);
+ Entity entity3 = new Entity(ENTITY_TYPE, ENTITY_ID3);
+ Entity entity4 = new Entity(otherEntityType, ENTITY_ID3);
+ EntityOwnershipListener listener = mock(EntityOwnershipListener.class);
+
+ // Register listener
+
+ shard.tell(new RegisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ // Register a couple candidates for the desired entity type and verify listener is notified.
+
+ shard.tell(new RegisterCandidateLocal(entity1), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity1, false, true, true));
+
+ shard.tell(new RegisterCandidateLocal(entity2), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity2, false, true, true));
+ reset(listener);
+
+ // Register another candidate for another entity type and verify listener is not notified.
+
+ shard.tell(new RegisterCandidateLocal(entity4), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
+ verify(listener, never()).ownershipChanged(ownershipChange(entity4));
+
+ // Register remote candidate for entity1
+
+ String remoteMemberName = "remoteMember";
+ writeNode(ENTITY_OWNERS_PATH, entityOwnersWithCandidate(ENTITY_TYPE, entity1.getId(), remoteMemberName),
+ shardDataTree);
+ verifyCommittedEntityCandidate(shard, ENTITY_TYPE, entity1.getId(), remoteMemberName);
+
+ // Unregister the local candidate for entity1 and verify listener is notified
+
+ shard.tell(new UnregisterCandidateLocal(entity1), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ verify(listener, timeout(5000)).ownershipChanged(ownershipChange(entity1, true, false, true));
+ reset(listener);
+
+ // Unregister the listener, add a candidate for entity3 and verify listener isn't notified
+
+ shard.tell(new UnregisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ shard.tell(new RegisterCandidateLocal(entity3), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ verifyOwner(shard, ENTITY_TYPE, entity3.getId(), LOCAL_MEMBER_NAME);
+ Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
+ verify(listener, never()).ownershipChanged(any(EntityOwnershipChange.class));
+
+ // Re-register the listener and verify it gets notified of currently owned entities
+
+ reset(listener);
+
+ shard.tell(new RegisterListenerLocal(listener, ENTITY_TYPE), kit.getRef());
+ kit.expectMsgClass(SuccessReply.class);
+
+ verify(listener, timeout(5000).times(2)).ownershipChanged(or(ownershipChange(entity2, false, true, true),
+ ownershipChange(entity3, false, true, true)));
+ Uninterruptibles.sleepUninterruptibly(300, TimeUnit.MILLISECONDS);
+ verify(listener, never()).ownershipChanged(ownershipChange(entity4));
+ verify(listener, times(1)).ownershipChanged(ownershipChange(entity1));
+ }
+
+ private static void commitModification(TestActorRef<EntityOwnershipShard> shard, NormalizedNode<?, ?> node,