InstrumentedDatastore implements Datastore, InstrumentedDatastoreReaderWriter { @Override public Batch newBatch() { return InstrumentedBatch.of(stats, delegate.newBatch()); } 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 newBatch() { when(datastore.newBatch()).thenReturn(batch); var instrumentedBatch = instrumentedDatastore.newBatch(); verify(datastore).newBatch(); verifyNoMoreInteractions(datastore); reset(datastore); testBatchWriter(instrumentedBatch, batch); instrumentedBatch.submit(); verify(batch).submit(); verifyNoMoreInteractions(batch); reset(batch); when(batch.getDatastore()).thenReturn(datastore); var instrumentedDatastore = instrumentedBatch.getDatastore(); assertThat(instrumentedDatastore, instanceOf(InstrumentedDatastore.class)); assertThat(((InstrumentedDatastore) instrumentedDatastore).delegate(), is(datastore)); verify(batch).getDatastore(); verifyNoMoreInteractions(batch); }