- static class TestShardBuilder extends EntityOwnershipShard.Builder {
- TestShardBuilder() {
- localMemberName("member-1").ownerSelectionStrategyConfig(
- EntityOwnerSelectionStrategyConfig.newBuilder().build());
- }
-
- private final AtomicReference<CountDownLatch> messageReceived = new AtomicReference<>();
- private final AtomicReference<Object> receivedMessage = new AtomicReference<>();
- private final AtomicReference<Class<?>> messageClass = new AtomicReference<>();
- private final AtomicReference<DataTree> dataTree = new AtomicReference<>();
-
- @Override
- public Props props() {
- verify();
- return Props.create(TestEntityOwnershipShard.class,this, messageClass, messageReceived,
- receivedMessage, dataTree);
- }
-
- @SuppressWarnings("unchecked")
- <T> T waitForShardMessage() {
- assertTrue("Message " + messageClass.get().getSimpleName() + " was not received",
- Uninterruptibles.awaitUninterruptibly(messageReceived.get(), 5, TimeUnit.SECONDS));
- assertEquals("Message type", messageClass.get(), receivedMessage.get().getClass());
- return (T) receivedMessage.get();
- }
-
- void expectShardMessage(Class<?> ofType) {
- messageReceived.set(new CountDownLatch(1));
- receivedMessage.set(null);
- messageClass.set(ofType);
- }
-
- void setDataTree(DataTree tree) {
- this.dataTree.set(tree);
- }
- }
-