Pass the coverage environment into tests using a nested set collected from the transitive closure instead of having a global environment in the configuration.

--
MOS_MIGRATED_REVID=128559756

15 files changed, 96 insertions, 45 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/OutputFileConfiguredTarget.java b/src/main/java/com/google/devtools/build/lib/analysis/OutputFileConfiguredTarget.java
index a7a56a0e6d..d7205a5bf5 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/OutputFileConfiguredTarget.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/OutputFileConfiguredTarget.java
@@ -19,6 +19,7 @@ import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.packages.OutputFile;
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesProvider;
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesProviderImpl;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 
 /**
@@ -82,6 +83,12 @@ public class OutputFileConfiguredTarget extends FileConfiguredTarget
         .getCoverageSupportFiles();
   }
 
+  @Override
+  public NestedSet<Pair<String, String>> getCoverageEnvironment() {
+    return getProvider(InstrumentedFilesProvider.class, InstrumentedFilesProviderImpl.EMPTY)
+        .getCoverageEnvironment();
+  }
+
   /**
    * Returns the corresponding provider from the generating rule, if it is non-null, or {@code
    * defaultValue} otherwise.
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java b/src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java
index eb9822f36e..ee56c4fbe1 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java
@@ -208,7 +208,7 @@ public final class RuleConfiguredTargetBuilder {
 
     TestEnvironmentProvider environmentProvider = findProvider(TestEnvironmentProvider.class);
     if (environmentProvider != null) {
-      testActionBuilder.setExtraEnv(environmentProvider.getEnvironment());
+      testActionBuilder.addExtraEnv(environmentProvider.getEnvironment());
     }
 
     final TestParams testParams =
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
index 01d09fc12d..8bcb611b84 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java
@@ -165,13 +165,6 @@ public final class BuildConfiguration {
       return implicitLabels;
     }
 
-    /**
-     * Returns a (key, value) mapping to insert into the subcommand environment for coverage.
-     */
-    public Map<String, String> getCoverageEnvironment() {
-      return ImmutableMap.<String, String>of();
-    }
-
     /*
      * Returns the command-line "Make" variable overrides.
      */
@@ -2086,18 +2079,6 @@ public final class BuildConfiguration {
   }
 
   /**
-   * Returns a (key, value) mapping to insert into the subcommand environment for coverage
-   * actions.
-   */
-  public Map<String, String> getCoverageEnvironment() {
-    Map<String, String> env = new HashMap<>();
-    for (Fragment fragment : fragments.values()) {
-      env.putAll(fragment.getCoverageEnvironment());
-    }
-    return env;
-  }
-
-  /**
    * Returns the default value for the specified "Make" variable for this
    * configuration.  Returns null if no value was found.
    */
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCommon.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCommon.java
index 684ac7a062..51c2b52c4a 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCommon.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcCommon.java
@@ -498,6 +498,7 @@ public final class CcCommon {
         : InstrumentedFilesCollector.collect(
             ruleContext, CppRuleClasses.INSTRUMENTATION_SPEC, CC_METADATA_COLLECTOR, files,
             CppHelper.getGcovFilesIfNeeded(ruleContext),
+            CppHelper.getCoverageEnvironmentIfNeeded(ruleContext),
             withBaselineCoverage);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchain.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchain.java
index 843711b39c..2c0f94150e 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchain.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchain.java
@@ -41,6 +41,7 @@ import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.packages.License;
 import com.google.devtools.build.lib.rules.RuleConfiguredTargetFactory;
 import com.google.devtools.build.lib.rules.cpp.FdoSupport.FdoException;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -186,6 +187,16 @@ public class CcToolchain implements RuleConfiguredTargetFactory {
     boolean supportsParamFiles = ruleContext.attributes().get("supports_param_files", BOOLEAN);
     boolean supportsHeaderParsing =
         ruleContext.attributes().get("supports_header_parsing", BOOLEAN);
+
+    NestedSetBuilder<Pair<String, String>> coverageEnvironment = NestedSetBuilder.compileOrder();
+
+    coverageEnvironment.add(Pair.of(
+        "COVERAGE_GCOV_PATH", cppConfiguration.getGcovExecutable().getPathString()));
+    if (cppConfiguration.getFdoInstrument() != null) {
+      coverageEnvironment.add(Pair.of(
+          "FDO_DIR", cppConfiguration.getFdoInstrument().getPathString()));
+    }
+
     CcToolchainProvider provider =
         new CcToolchainProvider(
             cppConfiguration,
@@ -206,7 +217,8 @@ public class CcToolchain implements RuleConfiguredTargetFactory {
             supportsParamFiles,
             supportsHeaderParsing,
             getBuildVariables(ruleContext),
-            getBuiltinIncludes(ruleContext));
+            getBuiltinIncludes(ruleContext),
+            coverageEnvironment.build());
     RuleConfiguredTargetBuilder builder =
         new RuleConfiguredTargetBuilder(ruleContext)
             .add(CcToolchainProvider.class, provider)
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchainProvider.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchainProvider.java
index 5b0847db7a..082d469810 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchainProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CcToolchainProvider.java
@@ -22,6 +22,7 @@ import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
 
@@ -57,7 +58,8 @@ public final class CcToolchainProvider implements TransitiveInfoProvider {
           false,
           false,
           ImmutableMap.<String, String>of(),
-          ImmutableList.<Artifact>of());
+          ImmutableList.<Artifact>of(),
+          NestedSetBuilder.<Pair<String, String>>emptySet(Order.COMPILE_ORDER));
 
   @Nullable private final CppConfiguration cppConfiguration;
   private final NestedSet<Artifact> crosstool;
@@ -78,6 +80,7 @@ public final class CcToolchainProvider implements TransitiveInfoProvider {
   private final boolean supportsHeaderParsing;
   private final Map<String, String> buildVariables;
   private final ImmutableList<Artifact> builtinIncludeFiles;
+  private final NestedSet<Pair<String, String>> coverageEnvironment;
 
   public CcToolchainProvider(
       @Nullable CppConfiguration cppConfiguration,
@@ -98,7 +101,8 @@ public final class CcToolchainProvider implements TransitiveInfoProvider {
       boolean supportsParamFiles,
       boolean supportsHeaderParsing,
       Map<String, String> buildVariables,
-      ImmutableList<Artifact> builtinIncludeFiles) {
+      ImmutableList<Artifact> builtinIncludeFiles,
+      NestedSet<Pair<String, String>> coverageEnvironment) {
     this.cppConfiguration = cppConfiguration;
     this.crosstool = Preconditions.checkNotNull(crosstool);
     this.crosstoolMiddleman = Preconditions.checkNotNull(crosstoolMiddleman);
@@ -118,6 +122,7 @@ public final class CcToolchainProvider implements TransitiveInfoProvider {
     this.supportsHeaderParsing = supportsHeaderParsing;
     this.buildVariables = buildVariables;
     this.builtinIncludeFiles = builtinIncludeFiles;
+    this.coverageEnvironment = coverageEnvironment;
   }
 
   /**
@@ -263,4 +268,11 @@ public final class CcToolchainProvider implements TransitiveInfoProvider {
   public ImmutableList<Artifact> getBuiltinIncludeFiles() {
     return builtinIncludeFiles;
   }
+
+  /**
+   * Returns the environment variables that need to be added to tests that collect code coverage.
+   */
+  public NestedSet<Pair<String, String>> getCoverageEnvironment() {
+    return coverageEnvironment;
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java
index 5924885fea..aa138ffa24 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java
@@ -1991,17 +1991,6 @@ public class CppConfiguration extends BuildConfiguration.Fragment {
   }
 
   @Override
-  public Map<String, String> getCoverageEnvironment() {
-    ImmutableMap.Builder<String, String> env = ImmutableMap.builder();
-    env.put("COVERAGE_GCOV_PATH", getGcovExecutable().getPathString());
-    PathFragment fdoInstrument = cppOptions.fdoInstrument;
-    if (fdoInstrument != null) {
-      env.put("FDO_DIR", fdoInstrument.getPathString());
-    }
-    return env.build();
-  }
-
-  @Override
   public String getOutputDirectoryName() {
     String lipoSuffix;
     if (getLipoMode() != LipoMode.OFF && !isAutoFdoLipo()) {
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppHelper.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppHelper.java
index 4c18120f97..6700f3ba54 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppHelper.java
@@ -43,6 +43,7 @@ import com.google.devtools.build.lib.rules.cpp.Link.LinkTargetType;
 import com.google.devtools.build.lib.shell.ShellUtils;
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.FileTypeSet;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.view.config.crosstool.CrosstoolConfig.LipoMode;
@@ -225,6 +226,15 @@ public class CppHelper {
         .getFdoSupport();
   }
 
+  public static NestedSet<Pair<String, String>> getCoverageEnvironmentIfNeeded(
+      RuleContext ruleContext) {
+    if (ruleContext.getConfiguration().isCodeCoverageEnabled()) {
+      return CppHelper.getToolchain(ruleContext).getCoverageEnvironment();
+    } else {
+      return NestedSetBuilder.emptySet(Order.COMPILE_ORDER);
+    }
+  }
+
   public static NestedSet<Artifact> getGcovFilesIfNeeded(RuleContext ruleContext) {
     if (ruleContext.getConfiguration().isCodeCoverageEnabled()) {
       return CppHelper.getToolchain(ruleContext).getCrosstool();
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCommon.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCommon.java
index e5eb9308f8..a87e5b4ba5 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCommon.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCommon.java
@@ -47,6 +47,7 @@ import com.google.devtools.build.lib.rules.test.InstrumentedFilesCollector.Local
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesProvider;
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.FileTypeSet;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -641,6 +642,7 @@ public class JavaCommon {
         JAVA_METADATA_COLLECTOR,
         filesToBuild,
         NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER),
+        NestedSetBuilder.<Pair<String, String>>emptySet(Order.STABLE_ORDER),
         /*withBaselineCoverage*/!TargetUtils.isTestRule(ruleContext.getTarget()));
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java b/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java
index c2396cac35..056dd97934 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java
@@ -91,6 +91,7 @@ import com.google.devtools.build.lib.rules.test.InstrumentedFilesCollector.Instr
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesCollector.LocalMetadataCollector;
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesProvider;
 import com.google.devtools.build.lib.util.FileTypeSet;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -887,6 +888,8 @@ public final class CompilationSupport {
         new ObjcCoverageMetadataCollector(),
         oFiles.build(),
         getGcovForObjectiveCIfNeeded(),
+        // The COVERAGE_GCOV_PATH environment variable is added in TestSupport#getExtraProviders()
+        NestedSetBuilder.<Pair<String, String>>emptySet(Order.COMPILE_ORDER),
         !TargetUtils.isTestRule(ruleContext.getTarget()));
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesCollector.java b/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesCollector.java
index d0357f16d7..1fb2fbfeea 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesCollector.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesCollector.java
@@ -30,6 +30,7 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.util.FileType;
 import com.google.devtools.build.lib.util.FileTypeSet;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 
 import java.util.ArrayList;
@@ -63,7 +64,9 @@ public final class InstrumentedFilesCollector {
       LocalMetadataCollector localMetadataCollector,
       Iterable<Artifact> rootFiles) {
     return collect(ruleContext, spec, localMetadataCollector, rootFiles,
-        NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER), false);
+        NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER),
+        NestedSetBuilder.<Pair<String, String>>emptySet(Order.STABLE_ORDER),
+        false);
   }
 
   /**
@@ -79,6 +82,7 @@ public final class InstrumentedFilesCollector {
       LocalMetadataCollector localMetadataCollector,
       Iterable<Artifact> rootFiles,
       NestedSet<Artifact> coverageSupportFiles,
+      NestedSet<Pair<String, String>> coverageEnvironment,
       boolean withBaselineCoverage) {
     Preconditions.checkNotNull(ruleContext);
     Preconditions.checkNotNull(spec);
@@ -94,6 +98,10 @@ public final class InstrumentedFilesCollector {
     NestedSetBuilder<Artifact> coverageSupportFilesBuilder =
         NestedSetBuilder.<Artifact>stableOrder()
             .addTransitive(coverageSupportFiles);
+    NestedSetBuilder<Pair<String, String>> coverageEnvironmentBuilder =
+        NestedSetBuilder.<Pair<String, String>>compileOrder()
+            .addTransitive(coverageEnvironment);
+
 
     // Transitive instrumentation data.
     for (TransitiveInfoCollection dep :
@@ -105,6 +113,7 @@ public final class InstrumentedFilesCollector {
         baselineCoverageInstrumentedFilesBuilder.addTransitive(
             provider.getBaselineCoverageInstrumentedFiles());
         coverageSupportFilesBuilder.addTransitive(provider.getCoverageSupportFiles());
+        coverageEnvironmentBuilder.addTransitive(provider.getCoverageEnvironment());
       }
     }
 
@@ -151,7 +160,8 @@ public final class InstrumentedFilesCollector {
         metadataFilesBuilder.build(),
         baselineCoverageFiles,
         baselineCoverageArtifacts,
-        coverageSupportFilesBuilder.build());
+        coverageSupportFilesBuilder.build(),
+        coverageEnvironmentBuilder.build());
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProvider.java b/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProvider.java
index 23e5e68733..82f787ec34 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProvider.java
@@ -17,6 +17,7 @@ package com.google.devtools.build.lib.rules.test;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.TransitiveInfoProvider;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
+import com.google.devtools.build.lib.util.Pair;
 
 /**
  * A provider of instrumented file sources and instrumentation metadata.
@@ -57,4 +58,9 @@ public interface InstrumentedFilesProvider extends TransitiveInfoProvider {
    * <p>They aren't mentioned in the instrumented files manifest.
    */
   NestedSet<Artifact> getCoverageSupportFiles();
+
+  /**
+   * Environment variables that need to be set for tests collecting code coverage.
+   */
+  NestedSet<Pair<String, String>> getCoverageEnvironment();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProviderImpl.java b/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProviderImpl.java
index 3397d2d35d..9be20a24ad 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProviderImpl.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/InstrumentedFilesProviderImpl.java
@@ -17,6 +17,7 @@ import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
+import com.google.devtools.build.lib.util.Pair;
 
 /**
  * An implementation class for the InstrumentedFilesProvider interface.
@@ -28,25 +29,29 @@ public final class InstrumentedFilesProviderImpl implements InstrumentedFilesPro
           NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER),
           NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER),
           NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER),
-          NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER));
+          NestedSetBuilder.<Artifact>emptySet(Order.STABLE_ORDER),
+          NestedSetBuilder.<Pair<String, String>>emptySet(Order.COMPILE_ORDER));
 
   private final NestedSet<Artifact> instrumentedFiles;
   private final NestedSet<Artifact> instrumentationMetadataFiles;
   private final NestedSet<Artifact> baselineCoverageFiles;
   private final NestedSet<Artifact> baselineCoverageArtifacts;
   private final NestedSet<Artifact> coverageSupportFiles;
+  private final NestedSet<Pair<String, String>> coverageEnvironment;
 
   public InstrumentedFilesProviderImpl(
       NestedSet<Artifact> instrumentedFiles,
       NestedSet<Artifact> instrumentationMetadataFiles,
       NestedSet<Artifact> baselineCoverageFiles,
       NestedSet<Artifact> baselineCoverageArtifacts,
-      NestedSet<Artifact> coverageSupportFiles) {
+      NestedSet<Artifact> coverageSupportFiles,
+      NestedSet<Pair<String, String>> coverageEnvironment) {
     this.instrumentedFiles = instrumentedFiles;
     this.instrumentationMetadataFiles = instrumentationMetadataFiles;
     this.baselineCoverageFiles = baselineCoverageFiles;
     this.baselineCoverageArtifacts = baselineCoverageArtifacts;
     this.coverageSupportFiles = coverageSupportFiles;
+    this.coverageEnvironment = coverageEnvironment;
   }
 
   @Override
@@ -73,4 +78,9 @@ public final class InstrumentedFilesProviderImpl implements InstrumentedFilesPro
   public NestedSet<Artifact> getCoverageSupportFiles() {
     return coverageSupportFiles;
   }
+
+  @Override
+  public NestedSet<Pair<String, String>> getCoverageEnvironment() {
+    return coverageEnvironment;
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/TestActionBuilder.java b/src/main/java/com/google/devtools/build/lib/rules/test/TestActionBuilder.java
index 6f3614f8e3..29be94d994 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/TestActionBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/TestActionBuilder.java
@@ -15,7 +15,6 @@
 package com.google.devtools.build.lib.rules.test;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Root;
@@ -33,12 +32,14 @@ import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.packages.TestSize;
 import com.google.devtools.build.lib.packages.TestTimeout;
 import com.google.devtools.build.lib.rules.test.TestProvider.TestParams;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.common.options.EnumConverter;
 
 import java.util.List;
 import java.util.Map;
+import java.util.TreeMap;
 
 import javax.annotation.Nullable;
 
@@ -57,7 +58,7 @@ public final class TestActionBuilder {
 
   public TestActionBuilder(RuleContext ruleContext) {
     this.ruleContext = ruleContext;
-    this.extraEnv = ImmutableMap.of();
+    this.extraEnv = new TreeMap<>();
   }
 
   /**
@@ -112,9 +113,8 @@ public final class TestActionBuilder {
     return this;
   }
 
-  public TestActionBuilder setExtraEnv(@Nullable Map<String, String> extraEnv) {
-    this.extraEnv = extraEnv == null
-        ? ImmutableMap.<String, String> of() : ImmutableMap.copyOf(extraEnv);
+  public TestActionBuilder addExtraEnv(Map<String, String> extraEnv) {
+    this.extraEnv.putAll(extraEnv);
     return this;
   }
 
@@ -197,6 +197,8 @@ public final class TestActionBuilder {
     final boolean collectCodeCoverage = config.isCodeCoverageEnabled()
         && instrumentedFiles != null;
 
+    TreeMap<String, String> testEnv = new TreeMap<>();
+
     TestTargetExecutionSettings executionSettings;
     if (collectCodeCoverage) {
       inputsBuilder.addTransitive(instrumentedFiles.getCoverageSupportFiles());
@@ -213,11 +215,16 @@ public final class TestActionBuilder {
       executionSettings = new TestTargetExecutionSettings(ruleContext, runfilesSupport,
           executable, instrumentedFileManifest, shards);
       inputsBuilder.add(instrumentedFileManifest);
+      for (Pair<String, String> coverageEnvEntry : instrumentedFiles.getCoverageEnvironment()) {
+        testEnv.put(coverageEnvEntry.getFirst(), coverageEnvEntry.getSecond());
+      }
     } else {
       executionSettings = new TestTargetExecutionSettings(ruleContext, runfilesSupport,
           executable, null, shards);
     }
 
+    testEnv.putAll(extraEnv);
+
     if (config.getRunUnder() != null) {
       Artifact runUnderExecutable = executionSettings.getRunUnderExecutable();
       if (runUnderExecutable != null) {
@@ -261,7 +268,7 @@ public final class TestActionBuilder {
             ruleContext.getActionOwner(), inputs, testRuntime,
             testLog, cacheStatus,
             coverageArtifact, microCoverageArtifact,
-            testProperties, extraEnv, executionSettings,
+            testProperties, testEnv, executionSettings,
             shard, run, config, ruleContext.getWorkspaceName()));
         results.add(cacheStatus);
       }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/TestEnvironmentProvider.java b/src/main/java/com/google/devtools/build/lib/rules/test/TestEnvironmentProvider.java
index 08b5bd97a0..0b8f01924b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/test/TestEnvironmentProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/test/TestEnvironmentProvider.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.rules.test;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.TransitiveInfoProvider;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
@@ -31,7 +32,7 @@ public final class TestEnvironmentProvider implements TransitiveInfoProvider {
    * Constructs a new provider with the given variable name to variable value mapping.
    */
   public TestEnvironmentProvider(ImmutableMap<String, String> environment) {
-    this.environment = environment;
+    this.environment = Preconditions.checkNotNull(environment);
   }
 
   /**