CountDataFunction implements DataFunction<Long> { @Override public <K, E> Long apply(DataStore<K, E> dataStore, QueryDescriptor query, RepositoryConfiguration repositoryConfiguration, List<E> selection) { if (selection == null) { log.error("Cannot calculate the count if the selection is null"); throw new InvalidArgumentException("Selection cannot be null"); } return ((Integer) selection.size()).longValue(); }  @Override Long apply(DataStore<K, E> dataStore, QueryDescriptor query,
                             RepositoryConfiguration repositoryConfiguration, List<E> selection);  }
@Test(expectedExceptions = InvalidArgumentException.class) public void testThatNullSelectionResultsInError() { final DataFunction function = new CountDataFunction(); function.apply(null, null, null, null); }
@Test public void testThatItReflectsTheSizeOfTheCollection() { final DataFunction<Long> function = new CountDataFunction(); final List<Object> selection = new LinkedList<>(); for (int i = 0; i < 10; i++) { final Long count = function.apply(null, null, null, selection); assertThat(count, Matchers.is((long) selection.size())); selection.add(new Object()); } }