Refactor cacheControl to more readable code (graphql-java#1492)

fcristovao · bbakerman · commit 8bae13f38c3c · 2019-04-09T15:22:03.000+10:00
diff --git a/src/main/java/graphql/cachecontrol/CacheControl.java b/src/main/java/graphql/cachecontrol/CacheControl.java
@@ -7,13 +7,11 @@
 import graphql.schema.DataFetchingEnvironment;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
 import static graphql.Assert.assertNotNull;
-import static graphql.Assert.assertTrue;
 import static graphql.util.FpKit.map;
 
 /**
@@ -31,6 +29,8 @@
 @PublicApi
 public class CacheControl {
 
+    public static final String CACHE_CONTROL_EXTENSION_KEY = "cacheControl";
+
     /**
      * If the scope is set to PRIVATE, this indicates anything under this path should only be cached per-user,
      * unless the value is overridden on a sub path. PUBLIC is the default and means anything under this path
@@ -170,26 +170,18 @@ public static CacheControl newCacheControl() {
      * @return a new execution result with the hints in the extensions map.
      */
     public ExecutionResult addTo(ExecutionResult executionResult) {
-        assertTrue(executionResult instanceof ExecutionResultImpl, "You must pass in an ExecutionResult based on graphql.ExecutionResultImpl");
-        Map<Object, Object> currentExtensions = executionResult.getExtensions();
-        if (currentExtensions == null) {
-            currentExtensions = Collections.emptyMap();
-        }
-        currentExtensions = new LinkedHashMap<>(currentExtensions);
-
-        putHintsInExtensionsMap(currentExtensions);
-
-        return ExecutionResultImpl.newExecutionResult().from(executionResult)
-                .extensions(currentExtensions).build();
+        return ExecutionResultImpl.newExecutionResult()
+                                  .from(executionResult)
+                                  .addExtension(CACHE_CONTROL_EXTENSION_KEY, hintsToCacheControlProperties())
+                                  .build();
     }
 
-    private void putHintsInExtensionsMap(Map<Object, Object> extensions) {
+    private Map<String, Object> hintsToCacheControlProperties() {
         List<Map<String, Object>> recordedHints = map(hints, Hint::toMap);
 
         Map<String, Object> cacheControl = new LinkedHashMap<>();
         cacheControl.put("version", 1);
         cacheControl.put("hints", recordedHints);
-
-        extensions.put("cacheControl", cacheControl);
+        return cacheControl;
     }
 }
