Merge pull request #100 from kelseyfrancis/prime-future

Allow priming with future values

Author: bbakerman

src/main/java/org/dataloader/DataLoader.java

@@ -685,21 +685,15 @@ public DataLoader<K, V> clearAll(BiConsumer<Void, Throwable> handler) {
     /**
      * Primes the cache with the given key and value.  Note this will only prime the future cache
      * and not the value store.  Use {@link ValueCache#set(Object, Object)} if you want
-     * o prime it with values before use
+     * to prime it with values before use
      *
      * @param key   the key
      * @param value the value
      *
      * @return the data loader for fluent coding
      */
     public DataLoader<K, V> prime(K key, V value) {
-        Object cacheKey = getCacheKey(key);
-        synchronized (this) {
-            if (!futureCache.containsKey(cacheKey)) {
-                futureCache.set(cacheKey, CompletableFuture.completedFuture(value));
-            }
-        }
-        return this;
+        return prime(key, CompletableFuture.completedFuture(value));
     }
 
     /**
@@ -711,9 +705,25 @@ public DataLoader<K, V> prime(K key, V value) {
      * @return the data loader for fluent coding
      */
     public DataLoader<K, V> prime(K key, Exception error) {
+        return prime(key, CompletableFutureKit.failedFuture(error));
+    }
+
+    /**
+     * Primes the cache with the given key and value.  Note this will only prime the future cache
+     * and not the value store.  Use {@link ValueCache#set(Object, Object)} if you want
+     * to prime it with values before use
+     *
+     * @param key   the key
+     * @param value the value
+     *
+     * @return the data loader for fluent coding
+     */
+    public DataLoader<K, V> prime(K key, CompletableFuture<V> value) {
         Object cacheKey = getCacheKey(key);
-        if (!futureCache.containsKey(cacheKey)) {
-            futureCache.set(cacheKey, CompletableFutureKit.failedFuture(error));
+        synchronized (this) {
+            if (!futureCache.containsKey(cacheKey)) {
+                futureCache.set(cacheKey, value);
+            }
         }
         return this;
     }

src/test/java/org/dataloader/DataLoaderTest.java

@@ -362,6 +362,24 @@ public void should_Allow_to_forcefully_prime_the_cache() throws ExecutionExcepti
         assertThat(loadCalls, equalTo(singletonList(singletonList("B"))));
     }
 
+    @Test
+    public void should_Allow_priming_the_cache_with_a_future() throws ExecutionException, InterruptedException {
+        List<Collection<String>> loadCalls = new ArrayList<>();
+        DataLoader<String, String> identityLoader = idLoader(new DataLoaderOptions(), loadCalls);
+
+        DataLoader<String, String> dlFluency = identityLoader.prime("A", CompletableFuture.completedFuture("A"));
+        assertThat(dlFluency, equalTo(identityLoader));
+
+        CompletableFuture<String> future1 = identityLoader.load("A");
+        CompletableFuture<String> future2 = identityLoader.load("B");
+        identityLoader.dispatch();
+
+        await().until(() -> future1.isDone() && future2.isDone());
+        assertThat(future1.get(), equalTo("A"));
+        assertThat(future2.get(), equalTo("B"));
+        assertThat(loadCalls, equalTo(singletonList(singletonList("B"))));
+    }
+
     @Test
     public void should_not_Cache_failed_fetches_on_complete_failure() {
         List<Collection<Integer>> loadCalls = new ArrayList<>();