-
- static class TestEntityOwnershipShard extends EntityOwnershipShard {
- private final AtomicReference<CountDownLatch> messageReceived;
- private final AtomicReference<Object> receivedMessage;
- private final AtomicReference<Class<?>> messageClass;
- private final AtomicReference<DataTree> dataTree;
-
- protected TestEntityOwnershipShard(final EntityOwnershipShard.Builder builder,
- final AtomicReference<Class<?>> messageClass, final AtomicReference<CountDownLatch> messageReceived,
- final AtomicReference<Object> receivedMessage, final AtomicReference<DataTree> dataTree) {
- super(builder);
- this.messageClass = messageClass;
- this.messageReceived = messageReceived;
- this.receivedMessage = receivedMessage;
- this.dataTree = dataTree;
- }
-
- @Override
- public void onReceiveCommand(final Object message) throws Exception {
- try {
- if(dataTree.get() != null && message instanceof GetShardDataTree) {
- sender().tell(dataTree.get(), self());
- } else {
- super.onReceiveCommand(message);
- }
- } finally {
- Class<?> expMsgClass = messageClass.get();
- if(expMsgClass != null && expMsgClass.equals(message.getClass())) {
- receivedMessage.set(message);
- messageReceived.get().countDown();
- }
- }
- }
- }