Extract parsing of the WORKSPACE file in two parts

The WORKSPACE file AST is now parsed as a separate SkyFunction
and this will be used to have multiple SkyValue for the same
WORKSPACE file, splitting the execution of the AST after load
statements to enable load statement of external dependencies
in the WORKSPACE file.

--
MOS_MIGRATED_REVID=112768897

14 files changed, 228 insertions, 109 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java b/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java
index 23c92dcd5d..8caec89db0 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/WorkspaceFactory.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.packages;
 
 import static com.google.devtools.build.lib.syntax.Runtime.NONE;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -54,8 +55,7 @@ public class WorkspaceFactory {
   public static final String BIND = "bind";
 
   private final LegacyBuilder builder;
-  private final StoredEventHandler localReporter;
-
+  
   private final Path installDir;
   private final Path workspaceDir;
   private final Mutability mutability;
@@ -93,7 +93,6 @@ public class WorkspaceFactory {
       @Nullable Path installDir,
       @Nullable Path workspaceDir) {
     this.builder = builder;
-    this.localReporter = new StoredEventHandler();
     this.mutability = mutability;
     this.installDir = installDir;
     this.workspaceDir = workspaceDir;
@@ -101,14 +100,6 @@ public class WorkspaceFactory {
     this.workspaceFunctions = createWorkspaceFunctions(ruleClassProvider);
   }
 
-  // State while parsing the WORKSPACE file.
-  // We store them so we can pause the parsing and load skylark imports from the
-  // WorkspaceFileFunction. Loading skylark imports require access to the .skyframe package
-  // which this package cannot depend on.
-  private BuildFileAST buildFileAST = null;
-  private Environment.Builder environmentBuilder = null;
-  private ParserInputSource source = null;
-
   /**
    * Parses the given WORKSPACE file without resolving skylark imports.
    *
@@ -116,44 +107,52 @@ public class WorkspaceFactory {
    * //src/tools/generate_workspace.</p>
    */
   public void parse(ParserInputSource source) throws InterruptedException, IOException {
+    parse(source, null);
+  }
+
+  @VisibleForTesting
+  public void parse(ParserInputSource source, @Nullable StoredEventHandler localReporter)
+      throws InterruptedException, IOException {
     // This method is split in 2 so WorkspaceFileFunction can call the two parts separately and
     // do the Skylark load imports in between. We can't load skylark imports from
     // generate_workspace at the moment because it doesn't have access to skyframe, but that's okay
     // because most people are just using it to resolve Maven dependencies.
-    parseWorkspaceFile(source);
-    execute();
-  }
-
-  /**
-   * Parses the WORKSPACE file, generating the AST.
-   * @throws IOException if parsing fails.
-   */
-  public void parseWorkspaceFile(ParserInputSource source) throws IOException {
-    this.source = source;
-    buildFileAST = BuildFileAST.parseBuildFile(source, localReporter, false);
+    if (localReporter == null) {
+      localReporter = new StoredEventHandler();
+    }
+    BuildFileAST buildFileAST = BuildFileAST.parseBuildFile(source, localReporter, false);
     if (buildFileAST.containsErrors()) {
-      environmentBuilder = null;
       throw new IOException("Failed to parse " + source.getPath());
     }
-    environmentBuilder = Environment.builder(mutability)
-        .setGlobals(Environment.BUILD)
-        .setEventHandler(localReporter);
+    execute(buildFileAST, null, localReporter);
   }
 
+
   /**
    * Actually runs through the AST, calling the functions in the WORKSPACE file and adding rules
    * to the //external package.
    */
-  public void execute() throws InterruptedException {
-    Preconditions.checkNotNull(environmentBuilder);
+  public void execute(BuildFileAST ast, Map<String, Extension> importedExtensions)
+      throws InterruptedException {
+    Preconditions.checkNotNull(ast);
+    Preconditions.checkNotNull(importedExtensions);
+    execute(ast, importedExtensions, new StoredEventHandler());
+  }
+  
+  private void execute(BuildFileAST ast, @Nullable Map<String, Extension> importedExtensions,
+      StoredEventHandler localReporter)
+      throws InterruptedException {
+    Environment.Builder environmentBuilder = Environment.builder(mutability)
+        .setGlobals(Environment.BUILD)
+        .setEventHandler(localReporter);
+    if (importedExtensions != null) {
+      environmentBuilder.setImportedExtensions(importedExtensions);
+    }
     Environment workspaceEnv = environmentBuilder.setLoadingPhase().build();
-    addWorkspaceFunctions(workspaceEnv);
-    if (!buildFileAST.exec(workspaceEnv, localReporter)) {
-      localReporter.handle(Event.error("Error evaluating WORKSPACE file " + source.getPath()));
+    addWorkspaceFunctions(workspaceEnv, localReporter);
+    if (!ast.exec(workspaceEnv, localReporter)) {
+      localReporter.handle(Event.error("Error evaluating WORKSPACE file"));
     }
-    environmentBuilder = null;
-    buildFileAST = null;
-    source = null;
 
     builder.addEvents(localReporter.getEvents());
     if (localReporter.hasErrors()) {
@@ -162,14 +161,6 @@ public class WorkspaceFactory {
     localReporter.clear();
   }
 
-  public BuildFileAST getBuildFileAST() {
-    return buildFileAST;
-  }
-
-  public void setImportedExtensions(Map<String, Extension> importedExtensions) {
-    environmentBuilder.setImportedExtensions(importedExtensions);
-  }
-
   // TODO(bazel-team): use @SkylarkSignature annotations on a BuiltinFunction.Factory
   // for signature + documentation of this and other functions in this file.
   private static BuiltinFunction newWorkspaceNameFunction() {
@@ -263,7 +254,7 @@ public class WorkspaceFactory {
     return mapBuilder.build();
   }
 
-  private void addWorkspaceFunctions(Environment workspaceEnv) {
+  private void addWorkspaceFunctions(Environment workspaceEnv, StoredEventHandler localReporter) {
     try {
       workspaceEnv.update("workspace", newWorkspaceNameFunction());
       for (Map.Entry<String, BaseFunction> function : workspaceFunctions.entrySet()) {
@@ -305,11 +296,4 @@ public class WorkspaceFactory {
   public static ClassObject newNativeModule(RuleClassProvider ruleClassProvider) {
     return newNativeModule(createWorkspaceFunctions(ruleClassProvider));
   }
-
-  /**
-   * @return a list of events that have occurred while parsing the WORKSPACE file.
-   */
-  public ImmutableList<Event> getEvents() {
-    return localReporter.getEvents();
-  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
index fcccc85791..277b627c9f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
@@ -96,6 +96,7 @@ public final class SkyFunctions {
   public static final SkyFunctionName WORKSPACE_FILE = SkyFunctionName.create("WORKSPACE_FILE");
   public static final SkyFunctionName COVERAGE_REPORT = SkyFunctionName.create("COVERAGE_REPORT");
   public static final SkyFunctionName REPOSITORY = SkyFunctionName.create("REPOSITORY");
+  public static final SkyFunctionName WORKSPACE_AST = SkyFunctionName.create("WORKSPACE_AST");
 
   public static Predicate<SkyKey> isSkyFunction(final SkyFunctionName functionName) {
     return new Predicate<SkyKey>() {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index a15ea8820d..7cddb5af16 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -378,6 +378,7 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
         configurationFactory));
     map.put(SkyFunctions.CONFIGURATION_FRAGMENT, new ConfigurationFragmentFunction(
         configurationFragments));
+    map.put(SkyFunctions.WORKSPACE_AST, new WorkspaceASTFunction(ruleClassProvider));
     map.put(
         SkyFunctions.WORKSPACE_FILE,
         new WorkspaceFileFunction(ruleClassProvider, pkgFactory, directories));
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTFunction.java
new file mode 100644
index 0000000000..d1077d2260
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTFunction.java
@@ -0,0 +1,83 @@
+// Copyright 2016 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.skyframe;
+
+import com.google.devtools.build.lib.packages.RuleClassProvider;
+import com.google.devtools.build.lib.syntax.BuildFileAST;
+import com.google.devtools.build.lib.syntax.ParserInputSource;
+import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.vfs.RootedPath;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyFunctionException;
+import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+
+import java.io.IOException;
+
+/**
+ * A SkyFunction to parse WORKSPACE files into a BuildFileAST.
+ */
+public class WorkspaceASTFunction implements SkyFunction {
+  private final RuleClassProvider ruleClassProvider;
+
+  public WorkspaceASTFunction(RuleClassProvider ruleClassProvider) {
+    this.ruleClassProvider = ruleClassProvider;
+  }
+
+  @Override
+  public SkyValue compute(SkyKey skyKey, Environment env)
+      throws InterruptedException, WorkspaceASTFunctionException {
+    RootedPath workspaceRoot = (RootedPath) skyKey.argument();
+    FileValue workspaceFileValue = (FileValue) env.getValue(FileValue.key(workspaceRoot));
+    if (workspaceFileValue == null) {
+      return null;
+    }
+
+    Path repoWorkspace = workspaceRoot.getRoot().getRelative(workspaceRoot.getRelativePath());
+    PathFragment pathFragment = new PathFragment("/DEFAULT.WORKSPACE");
+    try {
+      BuildFileAST ast = BuildFileAST.parseBuildFile(
+          ParserInputSource.create(ruleClassProvider.getDefaultWorkspaceFile(), pathFragment),
+          env.getListener(), false);
+      if (ast.containsErrors()) {
+        throw new WorkspaceASTFunctionException(
+            new IOException("Failed to parse default WORKSPACE file"), Transience.PERSISTENT);
+      }
+      if (workspaceFileValue.exists()) {
+        ast = BuildFileAST.parseBuildFile(
+            ParserInputSource.create(repoWorkspace), ast.getStatements(), env.getListener(), false);
+        if (ast.containsErrors()) {
+          throw new WorkspaceASTFunctionException(
+              new IOException("Failed to parse WORKSPACE file"), Transience.PERSISTENT);
+        }
+      }
+      return new WorkspaceASTValue(ast);
+    } catch (IOException ex) {
+      throw new WorkspaceASTFunctionException(ex, Transience.TRANSIENT);
+    }
+  }
+
+  private static final class WorkspaceASTFunctionException extends SkyFunctionException {
+    public WorkspaceASTFunctionException(Exception e, Transience transience) {
+      super(e, transience);
+    }
+  }
+  @Override
+  public String extractTag(SkyKey skyKey) {
+    return null;
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTValue.java
new file mode 100644
index 0000000000..022541d04f
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTValue.java
@@ -0,0 +1,42 @@
+// Copyright 2016 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.skyframe;
+
+import com.google.common.base.Preconditions;
+import com.google.devtools.build.lib.syntax.BuildFileAST;
+import com.google.devtools.build.lib.vfs.RootedPath;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+
+/**
+ * A SkyValue that stores the parsed WORKSPACE file as an AST.
+ */
+public class WorkspaceASTValue implements SkyValue {
+
+  private final BuildFileAST ast;
+
+  public WorkspaceASTValue(BuildFileAST ast) {
+    Preconditions.checkNotNull(ast);
+    this.ast = ast;
+  }
+
+  public BuildFileAST getAST() {
+    return ast;
+  }
+
+  public SkyKey key(RootedPath path) {
+    return new SkyKey(SkyFunctions.WORKSPACE_AST, path);
+  }
+}
+
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java
index 65060a29e1..dc2b0c2979 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunction.java
@@ -16,16 +16,14 @@ package com.google.devtools.build.lib.skyframe;
 
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.Package.LegacyBuilder;
 import com.google.devtools.build.lib.packages.PackageFactory;
 import com.google.devtools.build.lib.packages.RuleClassProvider;
 import com.google.devtools.build.lib.packages.WorkspaceFactory;
 import com.google.devtools.build.lib.skyframe.PackageFunction.PackageFunctionException;
+import com.google.devtools.build.lib.syntax.BuildFileAST;
 import com.google.devtools.build.lib.syntax.Mutability;
-import com.google.devtools.build.lib.syntax.ParserInputSource;
 import com.google.devtools.build.lib.vfs.Path;
-import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunctionException;
@@ -33,8 +31,6 @@ import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
-import java.io.IOException;
-
 /**
  * A SkyFunction to parse WORKSPACE files.
  */
@@ -56,11 +52,11 @@ public class WorkspaceFileFunction implements SkyFunction {
   @Override
   public SkyValue compute(SkyKey skyKey, Environment env) throws WorkspaceFileFunctionException,
       InterruptedException {
-    final Environment skyEnvironment = env;
 
     RootedPath workspaceRoot = (RootedPath) skyKey.argument();
-    FileValue workspaceFileValue = (FileValue) env.getValue(FileValue.key(workspaceRoot));
-    if (workspaceFileValue == null) {
+    WorkspaceASTValue workspaceASTValue =
+        (WorkspaceASTValue) env.getValue(new SkyKey(SkyFunctions.WORKSPACE_AST, workspaceRoot));
+    if (workspaceASTValue == null) {
       return null;
     }
 
@@ -78,60 +74,23 @@ public class WorkspaceFileFunction implements SkyFunction {
               directories.getEmbeddedBinariesRoot(),
               directories.getWorkspace());
       try {
-        PathFragment pathFragment = new PathFragment("/DEFAULT.WORKSPACE");
-        if (!parse(
-                ParserInputSource.create(ruleClassProvider.getDefaultWorkspaceFile(), pathFragment),
-                repoWorkspace, parser, skyEnvironment)) {
-          return null;
-        }
-        if (!workspaceFileValue.exists()) {
-          return new PackageValue(builder.build());
-        }
-
-        if (!parse(
-                ParserInputSource.create(repoWorkspace), repoWorkspace, parser, skyEnvironment)) {
+        BuildFileAST ast = workspaceASTValue.getAST();
+        PackageFunction.SkylarkImportResult importResult =
+            PackageFunction.fetchImportsFromBuildFile(
+                repoWorkspace, Label.EXTERNAL_PACKAGE_IDENTIFIER, ast, env, null);
+        if (importResult != null) {
+          parser.execute(ast, importResult.importMap);
+        } else {
           return null;
         }
       } catch (PackageFunctionException e) {
         throw new WorkspaceFileFunctionException(e, Transience.PERSISTENT);
-      } catch (IOException e) {
-        for (Event event : parser.getEvents()) {
-          env.getListener().handle(event);
-        }
-        throw new WorkspaceFileFunctionException(e, Transience.TRANSIENT);
       }
     }
 
     return new PackageValue(builder.build());
   }
 
-  private boolean loadSkylarkImports(Path repoWorkspace, WorkspaceFactory parser,
-      Environment skyEnvironment) throws PackageFunctionException, InterruptedException {
-    // Load skylark imports
-    PackageFunction.SkylarkImportResult importResult;
-    importResult = PackageFunction.fetchImportsFromBuildFile(repoWorkspace,
-            Label.EXTERNAL_PACKAGE_IDENTIFIER,
-            parser.getBuildFileAST(),
-            skyEnvironment,
-            null);
-    if (importResult == null) {
-      return false;
-    }
-    parser.setImportedExtensions(importResult.importMap);
-    return true;
-  }
-
-  private boolean parse(ParserInputSource source, Path repoWorkspace, WorkspaceFactory parser,
-      Environment skyEnvironment)
-      throws PackageFunctionException, InterruptedException, IOException {
-    parser.parseWorkspaceFile(source);
-    if (!loadSkylarkImports(repoWorkspace, parser, skyEnvironment)) {
-      return false;
-    }
-    parser.execute();
-    return true;
-  }
-
   @Override
   public String extractTag(SkyKey skyKey) {
     return null;
diff --git a/src/test/java/com/google/devtools/build/lib/packages/WorkspaceFactoryTest.java b/src/test/java/com/google/devtools/build/lib/packages/WorkspaceFactoryTest.java
index 5cdd3cb22e..13be11547d 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/WorkspaceFactoryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/WorkspaceFactoryTest.java
@@ -20,6 +20,7 @@ import static org.junit.Assert.fail;
 
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.StoredEventHandler;
 import com.google.devtools.build.lib.packages.Package.LegacyBuilder;
 import com.google.devtools.build.lib.syntax.Mutability;
 import com.google.devtools.build.lib.syntax.ParserInputSource;
@@ -63,14 +64,15 @@ public class WorkspaceFactoryTest {
             Mutability.create("test"),
             root,
             root);
+    StoredEventHandler localReporter = new StoredEventHandler();
     try {
-      factory.parse(ParserInputSource.create(workspaceFilePath));
+      factory.parse(ParserInputSource.create(workspaceFilePath), localReporter);
       fail("Parsing " + workspaceFilePath + " should have failed");
     } catch (IOException e) {
       assertThat(e.getMessage()).contains("Failed to parse " + workspaceFilePath);
     }
-    List<Event> events = factory.getEvents();
-    assertEquals(1, events.size());
+    List<Event> events = localReporter.getEvents();
+    assertEquals("Incorrect size for " + events, 1, events.size());
     assertThat(events.get(0).getMessage())
         .contains("The label must reference a file with extension '.bzl'");
   }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
index aae32548db..4870806cc5 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
@@ -117,6 +117,8 @@ public class ArtifactFunctionTest {
                 .put(SkyFunctions.PACKAGE,
                     new PackageFunction(null, null, null, null, null, null, null))
                 .put(SkyFunctions.PACKAGE_LOOKUP, new PackageLookupFunction(null))
+                .put(SkyFunctions.WORKSPACE_AST,
+                    new WorkspaceASTFunction(TestRuleClassProvider.getRuleClassProvider()))
                 .put(SkyFunctions.WORKSPACE_FILE,
                     new WorkspaceFileFunction(TestRuleClassProvider.getRuleClassProvider(),
                         new PackageFactory(TestRuleClassProvider.getRuleClassProvider()),
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/FileFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/FileFunctionTest.java
index e3f1c9e478..7780f7fc24 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/FileFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/FileFunctionTest.java
@@ -140,6 +140,8 @@ public class FileFunctionTest {
                 .put(SkyFunctions.PACKAGE_LOOKUP,
                     new PackageLookupFunction(new AtomicReference<>(
                         ImmutableSet.<PackageIdentifier>of())))
+                .put(SkyFunctions.WORKSPACE_AST,
+                    new WorkspaceASTFunction(TestRuleClassProvider.getRuleClassProvider()))
                 .put(SkyFunctions.WORKSPACE_FILE,
                     new WorkspaceFileFunction(TestRuleClassProvider.getRuleClassProvider(),
                         new PackageFactory(TestRuleClassProvider.getRuleClassProvider()),
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java
index 01111afb26..3b0380e15f 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java
@@ -113,6 +113,8 @@ public class FilesystemValueCheckerTest {
         new PackageFunction(null, null, null, null, null, null, null));
     skyFunctions.put(SkyFunctions.PACKAGE_LOOKUP,
         new PackageLookupFunction(new AtomicReference<>(ImmutableSet.<PackageIdentifier>of())));
+    skyFunctions.put(SkyFunctions.WORKSPACE_AST,
+        new WorkspaceASTFunction(TestRuleClassProvider.getRuleClassProvider()));
     skyFunctions.put(SkyFunctions.WORKSPACE_FILE,
         new WorkspaceFileFunction(TestRuleClassProvider.getRuleClassProvider(),
             new PackageFactory(TestRuleClassProvider.getRuleClassProvider()),
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/PackageLookupFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/PackageLookupFunctionTest.java
index a27cb61bb3..4195a3db15 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/PackageLookupFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/PackageLookupFunctionTest.java
@@ -88,6 +88,8 @@ public class PackageLookupFunctionTest extends FoundationTestCase {
     skyFunctions.put(SkyFunctions.BLACKLISTED_PACKAGE_PREFIXES,
         new BlacklistedPackagePrefixesFunction());
     RuleClassProvider ruleClassProvider = TestRuleClassProvider.getRuleClassProvider();
+    skyFunctions.put(SkyFunctions.WORKSPACE_AST,
+        new WorkspaceASTFunction(TestRuleClassProvider.getRuleClassProvider()));
     skyFunctions.put(
         SkyFunctions.WORKSPACE_FILE,
         new WorkspaceFileFunction(
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java
index 5c32da996f..62a375f460 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/RecursiveFilesystemTraversalFunctionTest.java
@@ -108,6 +108,8 @@ public final class RecursiveFilesystemTraversalFunctionTest extends FoundationTe
         new PackageFunction(null, null, null, null, null, null, null));
     skyFunctions.put(SkyFunctions.PACKAGE_LOOKUP,
         new PackageLookupFunction(deletedPackages));
+    skyFunctions.put(SkyFunctions.WORKSPACE_AST,
+        new WorkspaceASTFunction(TestRuleClassProvider.getRuleClassProvider()));
     skyFunctions.put(SkyFunctions.WORKSPACE_FILE,
         new WorkspaceFileFunction(TestRuleClassProvider.getRuleClassProvider(),
             new PackageFactory(TestRuleClassProvider.getRuleClassProvider()),
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
index 7f24a34cb0..995ece51db 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
@@ -162,6 +162,8 @@ public abstract class TimestampBuilderTestCase extends FoundationTestCase {
                 .put(SkyFunctions.PACKAGE,
                     new PackageFunction(null, null, null, null, null, null, null))
                 .put(SkyFunctions.PACKAGE_LOOKUP, new PackageLookupFunction(null))
+                .put(SkyFunctions.WORKSPACE_AST,
+                    new WorkspaceASTFunction(TestRuleClassProvider.getRuleClassProvider()))
                 .put(SkyFunctions.WORKSPACE_FILE,
                     new WorkspaceFileFunction(TestRuleClassProvider.getRuleClassProvider(),
                         new PackageFactory(TestRuleClassProvider.getRuleClassProvider()),
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunctionTest.java
index 91afbfaae7..71c3cf296f 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/WorkspaceFileFunctionTest.java
@@ -32,14 +32,20 @@ import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
 
+import org.hamcrest.BaseMatcher;
+import org.hamcrest.Description;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Matchers;
 import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 
 import java.io.IOException;
 
@@ -50,6 +56,7 @@ import java.io.IOException;
 public class WorkspaceFileFunctionTest extends BuildViewTestCase {
 
   private WorkspaceFileFunction skyFunc;
+  private WorkspaceASTFunction astSkyFunc;
   private FakeFileValue fakeWorkspaceFileValue;
 
   private static class FakeFileValue extends FileValue {
@@ -94,13 +101,13 @@ public class WorkspaceFileFunctionTest extends BuildViewTestCase {
   @Before
   public final void setUp() throws Exception {
     ConfiguredRuleClassProvider ruleClassProvider = TestRuleClassProvider.getRuleClassProvider();
-    skyFunc =
+    skyFunc = 
         new WorkspaceFileFunction(
             ruleClassProvider,
-            new PackageFactory(
-                TestRuleClassProvider.getRuleClassProvider(),
+            new PackageFactory(ruleClassProvider,
                 new BazelRulesModule().getPackageEnvironmentExtension()),
             directories);
+    astSkyFunc = new WorkspaceASTFunction(ruleClassProvider);
     fakeWorkspaceFileValue = new FakeFileValue();
   }
 
@@ -115,9 +122,37 @@ public class WorkspaceFileFunctionTest extends BuildViewTestCase {
         workspacePath.getParentDirectory(), new PathFragment(workspacePath.getBaseName()));
   }
 
+  // Dummy harmcrest matcher that match the function name of a skykey
+  private static class SkyKeyMatchers extends BaseMatcher<SkyKey> {
+    private final SkyFunctionName functionName;
+    
+    public SkyKeyMatchers(SkyFunctionName functionName) {
+      this.functionName = functionName;
+    }
+    @Override
+    public boolean matches(Object item) {
+      if (item instanceof SkyKey) {
+        return ((SkyKey) item).functionName().equals(functionName);
+      }
+      return false;
+    }
+    
+    @Override
+    public void describeTo(Description description) {}
+  }
+  
   private SkyFunction.Environment getEnv() {
     SkyFunction.Environment env = Mockito.mock(SkyFunction.Environment.class);
-    Mockito.when(env.getValue(Matchers.<SkyKey>any())).thenReturn(fakeWorkspaceFileValue);
+    Mockito.when(env.getValue(Matchers.argThat(new SkyKeyMatchers(SkyFunctions.FILE))))
+        .thenReturn(fakeWorkspaceFileValue);
+    Mockito.when(env.getValue(Matchers.argThat(new SkyKeyMatchers(SkyFunctions.WORKSPACE_AST))))
+    .then(new Answer<SkyValue>() {
+      @Override
+      public SkyValue answer(InvocationOnMock invocation) throws Throwable {
+        SkyKey key = (SkyKey) invocation.getArguments()[0];
+        return astSkyFunc.compute(key, getEnv());
+      }
+    });
     return env;
   }