Allow injection of custom SkyFunctions into GraphTester.

PiperOrigin-RevId: 199126212

1 files changed, 10 insertions, 3 deletions

diff --git a/src/test/java/com/google/devtools/build/skyframe/GraphTester.java b/src/test/java/com/google/devtools/build/skyframe/GraphTester.java
index 8901ae47f9..9b129df70e 100644
--- a/src/test/java/com/google/devtools/build/skyframe/GraphTester.java
+++ b/src/test/java/com/google/devtools/build/skyframe/GraphTester.java
@@ -50,12 +50,15 @@ import javax.annotation.Nullable;
 public class GraphTester {
 
   public static final SkyFunctionName NODE_TYPE = SkyFunctionName.FOR_TESTING;
-  private final ImmutableMap<SkyFunctionName, ? extends SkyFunction> functionMap =
-      ImmutableMap.of(GraphTester.NODE_TYPE, new DelegatingFunction());
+  private final Map<SkyFunctionName, SkyFunction> functionMap = new HashMap<>();
 
   private final Map<SkyKey, TestFunction> values = new HashMap<>();
   private final Set<SkyKey> modifiedValues = new LinkedHashSet<>();
 
+  public GraphTester() {
+    functionMap.put(NODE_TYPE, new DelegatingFunction());
+  }
+
   public TestFunction getOrCreate(String name) {
     return getOrCreate(skyKey(name));
   }
@@ -309,7 +312,11 @@ public class GraphTester {
   }
 
   public ImmutableMap<SkyFunctionName, ? extends SkyFunction> getSkyFunctionMap() {
-    return functionMap;
+    return ImmutableMap.copyOf(functionMap);
+  }
+
+  public void putSkyFunction(SkyFunctionName functionName, SkyFunction function) {
+    functionMap.put(functionName, function);
   }
 
   /**