Allow injection of custom SkyFunctions into GraphTester. PiperOrigin-RevId: 199126212
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 8901ae4..9b129df 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 @@ 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 ImmutableMap<SkyFunctionName, ? extends SkyFunction> getSkyFunctionMap() { - return functionMap; + return ImmutableMap.copyOf(functionMap); + } + + public void putSkyFunction(SkyFunctionName functionName, SkyFunction function) { + functionMap.put(functionName, function); } /**