DatastoreStorage implements Closeable { Map<WorkflowInstance, RunState> readActiveStates() throws IOException { var timeout = CompletableFuture.runAsync(() -> {}, delayedExecutor(30, SECONDS)); var instances = listActiveInstances0(timeout); var states = gatherIO(Lists.partition(List.copyOf(instances), MAX_NUMBER_OF_ENTITIES_IN_ONE_BATCH_READ).stream() .map(batch -> asyncIO(() -> readRunStateBatch(batch))) .collect(toList()), timeout) .stream() .flatMap(Collection::stream) .collect(toMap(RunState::workflowInstance, Function.identity())); timeout.cancel(true); return states; } DatastoreStorage(CheckedDatastore datastore); @VisibleForTesting  DatastoreStorage(CheckedDatastore datastore,
                   Function<CheckedDatastoreTransaction, DatastoreStorageTransaction> storageTransactionFactory,
                   ExecutorService executor, Logger log); @Override void close(); void updateNextNaturalTrigger(WorkflowId workflowId, TriggerInstantSpec triggerSpec); Map<WorkflowId, Workflow> workflows(); Map<WorkflowId, Workflow> workflows(Set<WorkflowId> workflowIds); List<Workflow> workflows(String componentId); T runInTransactionWithRetries(TransactionFunction<T, E> f);  }
@Test public void readActiveStatesShouldPropagateIOException() throws Exception { var cause = new IOException("foobar"); doThrow(cause).when(datastore).query(any()); var exception = assertThrows(IOException.class, () -> storage.readActiveStates()); assertThat(exception, is(cause)); }