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
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 23c92dc..8caec89 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 @@
 
 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 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 @@
       @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 @@
     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 @@
    * //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);
-    Environment workspaceEnv = environmentBuilder.setLoadingPhase().build();
-    addWorkspaceFunctions(workspaceEnv);
-    if (!buildFileAST.exec(workspaceEnv, localReporter)) {
-      localReporter.handle(Event.error("Error evaluating WORKSPACE file " + source.getPath()));
+  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);
     }
-    environmentBuilder = null;
-    buildFileAST = null;
-    source = null;
+    Environment workspaceEnv = environmentBuilder.setLoadingPhase().build();
+    addWorkspaceFunctions(workspaceEnv, localReporter);
+    if (!ast.exec(workspaceEnv, localReporter)) {
+      localReporter.handle(Event.error("Error evaluating WORKSPACE file"));
+    }
 
     builder.addEvents(localReporter.getEvents());
     if (localReporter.hasErrors()) {
@@ -162,14 +161,6 @@
     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 @@
     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 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 fcccc85..277b627 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 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 a15ea88..7cddb5a 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 @@
         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 0000000..d1077d2
--- /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 0000000..022541d
--- /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 65060a2..dc2b0c2 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 @@
 
 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.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
-import java.io.IOException;
-
 /**
  * A SkyFunction to parse WORKSPACE files.
  */
@@ -56,11 +52,11 @@
   @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 @@
               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 5cdd3cb..13be115 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 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 @@
             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 aae3254..4870806 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 @@
                 .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 e3f1c9e..7780f7f 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 @@
                 .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 01111af..3b0380e 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 @@
         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 a27cb61..4195a3d 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 @@
     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 5c32da9..62a375f 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 @@
         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 7f24a34..995ece5 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 @@
                 .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 91afbfa..71c3cf2 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.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 @@
 public class WorkspaceFileFunctionTest extends BuildViewTestCase {
 
   private WorkspaceFileFunction skyFunc;
+  private WorkspaceASTFunction astSkyFunc;
   private FakeFileValue fakeWorkspaceFileValue;
 
   private static class FakeFileValue extends FileValue {
@@ -94,13 +101,13 @@
   @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 @@
         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;
   }