diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index fe9c59a..05abf42 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -685,7 +685,7 @@ public DataLoader clearAll(BiConsumer 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 @@ -693,13 +693,7 @@ public DataLoader clearAll(BiConsumer handler) { * @return the data loader for fluent coding */ public DataLoader 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 prime(K key, V value) { * @return the data loader for fluent coding */ public DataLoader 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 prime(K key, CompletableFuture 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; } diff --git a/src/test/java/org/dataloader/DataLoaderTest.java b/src/test/java/org/dataloader/DataLoaderTest.java index 08d5b44..63a834d 100644 --- a/src/test/java/org/dataloader/DataLoaderTest.java +++ b/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> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + DataLoader dlFluency = identityLoader.prime("A", CompletableFuture.completedFuture("A")); + assertThat(dlFluency, equalTo(identityLoader)); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture 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> loadCalls = new ArrayList<>();