InstrumentedDatastore implements Datastore, InstrumentedDatastoreReaderWriter { @Override public <T> T runInTransaction(TransactionCallable<T> transactionCallable) { return delegate.runInTransaction(rw -> transactionCallable.run(InstrumentedDatastoreReaderWriter.of(stats, rw))); } private  InstrumentedDatastore(Datastore delegate, Stats stats); @Override Transaction newTransaction(TransactionOptions transactionOptions); @Override Transaction newTransaction(); @Override T runInTransaction(TransactionCallable<T> transactionCallable); @Override T runInTransaction(TransactionCallable<T> transactionCallable, TransactionOptions transactionOptions); @Override Batch newBatch(); @Override Entity get(Key key, ReadOption... readOptions); @Override Iterator<Entity> get(Iterable<Key> keys, ReadOption... readOptions); @Override List<Entity> fetch(Iterable<Key> keys, ReadOption... readOptions); @Override QueryResults<T> run(Query<T> query, ReadOption... readOptions); @Override Key allocateId(IncompleteKey key); @Override List<Key> allocateId(IncompleteKey... keys); @Override Entity add(FullEntity<?> entity); @Override List<Entity> add(FullEntity<?>... entities); @Override void update(Entity... entities); @Override Entity put(FullEntity<?> entity); @Override List<Entity> put(FullEntity<?>... entities); @Override void delete(Key... keys); @Override KeyFactory newKeyFactory(); @Override DatastoreOptions getOptions(); @Override Stats stats(); @Override DatastoreReaderWriter readerWriter(); Datastore delegate(); static InstrumentedDatastore of(Datastore delegate, Stats stats);  }
@Test public void runInTransaction() { when(datastore.runInTransaction(any())).then(a -> a.<TransactionCallable<String>>getArgument(0) .run(transaction)); var foobar = instrumentedDatastore.runInTransaction(tx -> { testReaderWriter(tx, transaction); return "foobar"; }); verify(datastore).runInTransaction(any()); verifyNoMoreInteractions(datastore); assertThat(foobar, is("foobar")); }