+
+ public static class CapturingShardDataTreeCohort extends ShardDataTreeCohort {
+ private volatile ShardDataTreeCohort delegate;
+ private FutureCallback<Empty> canCommit;
+ private FutureCallback<DataTreeCandidate> preCommit;
+ private FutureCallback<UnsignedLong> commit;
+
+ public void setDelegate(final ShardDataTreeCohort delegate) {
+ this.delegate = delegate;
+ }
+
+ public FutureCallback<Empty> getCanCommit() {
+ assertNotNull("canCommit was not invoked", canCommit);
+ return canCommit;
+ }
+
+ public FutureCallback<DataTreeCandidate> getPreCommit() {
+ assertNotNull("preCommit was not invoked", preCommit);
+ return preCommit;
+ }
+
+ public FutureCallback<UnsignedLong> getCommit() {
+ assertNotNull("commit was not invoked", commit);
+ return commit;
+ }
+
+ @Override
+ public TransactionIdentifier getIdentifier() {
+ return delegate.getIdentifier();
+ }
+
+ @Override
+ DataTreeCandidateTip getCandidate() {
+ return delegate.getCandidate();
+ }
+
+ @Override
+ DataTreeModification getDataTreeModification() {
+ return delegate.getDataTreeModification();
+ }
+
+ @Override
+ public void canCommit(final FutureCallback<Empty> callback) {
+ canCommit = mockFutureCallback(callback);
+ delegate.canCommit(canCommit);
+ }
+
+ @Override
+ public void preCommit(final FutureCallback<DataTreeCandidate> callback) {
+ preCommit = mockFutureCallback(callback);
+ delegate.preCommit(preCommit);
+ }
+
+ @Override
+ public void commit(final FutureCallback<UnsignedLong> callback) {
+ commit = mockFutureCallback(callback);
+ delegate.commit(commit);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> FutureCallback<T> mockFutureCallback(final FutureCallback<T> actual) {
+ FutureCallback<T> mock = mock(FutureCallback.class);
+ doAnswer(invocation -> {
+ actual.onFailure(invocation.getArgument(0));
+ return null;
+ }).when(mock).onFailure(any(Throwable.class));
+
+ doAnswer(invocation -> {
+ actual.onSuccess(invocation.getArgument(0));
+ return null;
+ }).when(mock).onSuccess((T) nullable(Object.class));
+
+ return mock;
+ }
+
+ @Override
+ public void abort(final FutureCallback<Empty> callback) {
+ delegate.abort(callback);
+ }
+
+ @Override
+ public boolean isFailed() {
+ return delegate.isFailed();
+ }
+
+ @Override
+ public State getState() {
+ return delegate.getState();
+ }
+
+ @Override
+ Optional<SortedSet<String>> getParticipatingShardNames() {
+ return delegate.getParticipatingShardNames();
+ }
+ }