+
+ // Add a new replica
+
+ JavaTestKit mockShardLeaderKit = new JavaTestKit(getSystem());
+
+ TestShardManager shardManagerInstance = shardManager.underlyingActor();
+ shardManagerInstance.setMessageInterceptor(newFindPrimaryInterceptor(mockShardLeaderKit.getRef()));
+
+ shardManager.tell(new AddShardReplica("astronauts"), kit.getRef());
+ mockShardLeaderKit.expectMsgClass(AddServer.class);
+ mockShardLeaderKit.reply(new AddServerReply(ServerChangeStatus.OK, ""));
+ kit.expectMsgClass(Status.Success.class);
+ waitForShardInitialized(shardManager, "astronauts", kit);
+
+ // Send another GetSnapshot and verify
+
+ shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
+ datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
+
+ assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"), Sets.newHashSet(
+ Lists.transform(datastoreSnapshot.getShardSnapshots(), shardNameTransformer)));
+
+ byte[] snapshotBytes = datastoreSnapshot.getShardManagerSnapshot();
+ assertNotNull("Expected ShardManagerSnapshot", snapshotBytes);
+ ShardManagerSnapshot snapshot = SerializationUtils.deserialize(snapshotBytes);
+ assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"),
+ Sets.newHashSet(snapshot.getShardList()));
+
+ LOG.info("testGetSnapshot ending");
+ }
+
+ @Test
+ public void testRestoreFromSnapshot() throws Throwable {
+ LOG.info("testRestoreFromSnapshot starting");
+
+ JavaTestKit kit = new JavaTestKit(getSystem());
+
+ MockConfiguration mockConfig = new MockConfiguration(ImmutableMap.<String, List<String>>builder().
+ put("shard1", Collections.<String>emptyList()).
+ put("shard2", Collections.<String>emptyList()).
+ put("astronauts", Collections.<String>emptyList()).build());
+
+
+ ShardManagerSnapshot snapshot = new ShardManagerSnapshot(Arrays.asList("shard1", "shard2", "astronauts"));
+ DatastoreSnapshot restoreFromSnapshot = new DatastoreSnapshot(shardMrgIDSuffix,
+ SerializationUtils.serialize(snapshot), Collections.<ShardSnapshot>emptyList());
+ TestActorRef<TestShardManager> shardManager = actorFactory.createTestActor(newTestShardMgrBuilder(mockConfig).
+ restoreFromSnapshot(restoreFromSnapshot).props().withDispatcher(Dispatchers.DefaultDispatcherId()));
+
+ shardManager.underlyingActor().waitForRecoveryComplete();
+
+ shardManager.tell(new UpdateSchemaContext(TestModel.createTestContext()), ActorRef.noSender());
+
+ waitForShardInitialized(shardManager, "shard1", kit);
+ waitForShardInitialized(shardManager, "shard2", kit);
+ waitForShardInitialized(shardManager, "astronauts", kit);
+
+ shardManager.tell(GetSnapshot.INSTANCE, kit.getRef());
+
+ DatastoreSnapshot datastoreSnapshot = expectMsgClassOrFailure(DatastoreSnapshot.class, kit, "GetSnapshot");
+
+ assertEquals("getType", shardMrgIDSuffix, datastoreSnapshot.getType());
+
+ byte[] snapshotBytes = datastoreSnapshot.getShardManagerSnapshot();
+ assertNotNull("Expected ShardManagerSnapshot", snapshotBytes);
+ snapshot = SerializationUtils.deserialize(snapshotBytes);
+ assertEquals("Shard names", Sets.newHashSet("shard1", "shard2", "astronauts"),
+ Sets.newHashSet(snapshot.getShardList()));
+
+ LOG.info("testRestoreFromSnapshot ending");