Refactors environment variables handling to make it testable. (#9586)

Author: AlexeyKuznetsov-DD

components/environment/src/main/java/datadog/environment/EnvironmentVariables.java

@@ -16,6 +16,19 @@
 public final class EnvironmentVariables {
   private EnvironmentVariables() {}
 
+  public static class EnvironmentVariablesProvider {
+    public String get(String name) {
+      return System.getenv(name);
+    }
+
+    public Map<String, String> getAll() {
+      return System.getenv();
+    }
+  }
+
+  // Make it accessible from tests.
+  public static EnvironmentVariablesProvider provider = new EnvironmentVariablesProvider();
+
   /**
    * Gets an environment variable value.
    *
@@ -41,7 +54,7 @@ public static String getOrDefault(String name, String defaultValue) {
       return defaultValue;
     }
     try {
-      String value = System.getenv(name);
+      String value = provider.get(name);
       return value == null ? defaultValue : value;
     } catch (SecurityException e) {
       return defaultValue;
@@ -56,7 +69,7 @@ public static String getOrDefault(String name, String defaultValue) {
    */
   public static Map<String, String> getAll() {
     try {
-      return unmodifiableMap(new HashMap<>(System.getenv()));
+      return unmodifiableMap(new HashMap<>(provider.getAll()));
     } catch (SecurityException e) {
       return emptyMap();
     }

dd-java-agent/agent-builder/src/test/groovy/datadog/trace/agent/test/DefaultInstrumenterForkedTest.groovy

@@ -1,6 +1,6 @@
 package datadog.trace.agent.test
 
-
+import datadog.environment.EnvironmentVariables
 import datadog.trace.agent.tooling.InstrumenterModule
 import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers
 import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade
@@ -118,7 +118,7 @@ class DefaultInstrumenterForkedTest extends DDSpecification {
     def target = new TestDefaultInstrumenter(name, altName)
 
     then:
-    System.getenv("DD_INTEGRATION_${value}_ENABLED") == "true"
+    EnvironmentVariables.get("DD_INTEGRATION_${value}_ENABLED") == "true"
     target.enabled == enabled
 
     where:

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/ci/CIProviderInfoFactoryTest.groovy

@@ -2,9 +2,7 @@ package datadog.trace.civisibility.ci
 
 import datadog.trace.api.Config
 import datadog.trace.civisibility.ci.env.CiEnvironmentImpl
-import org.junit.Rule
-import org.junit.contrib.java.lang.system.EnvironmentVariables
-import org.junit.contrib.java.lang.system.RestoreSystemProperties
+import datadog.trace.test.util.ControllableEnvironmentVariables
 import spock.lang.Specification
 
 import java.nio.file.Paths
@@ -21,25 +19,18 @@ import static datadog.trace.civisibility.ci.JenkinsInfo.JENKINS
 import static datadog.trace.civisibility.ci.TravisInfo.TRAVIS
 
 class CIProviderInfoFactoryTest extends Specification {
-  @Rule
-  public final EnvironmentVariables environmentVariables = new EnvironmentVariables()
+  protected ControllableEnvironmentVariables env = ControllableEnvironmentVariables.setup()
 
-  @Rule
-  public final RestoreSystemProperties restoreSystemProperties = new RestoreSystemProperties()
-
-  def setup() {
-    // Clear all environment variables to avoid clashes between
-    // real CI/Git environment variables and the spec CI/Git
-    // environment variables.
-    environmentVariables.clear(System.getenv().keySet() as String[])
+  void cleanup() {
+    env.clear()
   }
 
   def "test correct info is selected"() {
     setup:
-    environmentVariables.set(ciKeySelector, "true")
+    env.set(ciKeySelector, "true")
 
     when:
-    def ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), new CiEnvironmentImpl(System.getenv()))
+    def ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), new CiEnvironmentImpl(env.getAll()))
     def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(Paths.get("").toAbsolutePath())
 
     then:

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/ci/CITagsProviderTest.groovy

@@ -8,9 +8,7 @@ import datadog.trace.civisibility.ci.env.CiEnvironmentImpl
 import datadog.trace.civisibility.git.CILocalGitInfoBuilder
 import datadog.trace.civisibility.git.CIProviderGitInfoBuilder
 import datadog.trace.civisibility.git.tree.GitClient
-import org.junit.Rule
-import org.junit.contrib.java.lang.system.EnvironmentVariables
-import org.junit.contrib.java.lang.system.RestoreSystemProperties
+import datadog.trace.test.util.ControllableEnvironmentVariables
 import spock.lang.Specification
 
 import java.nio.file.Path
@@ -19,36 +17,28 @@ import java.nio.file.Paths
 import static datadog.trace.util.ConfigStrings.propertyNameToEnvironmentVariableName
 
 abstract class CITagsProviderTest extends Specification {
-
   static final CI_WORKSPACE_PATH_FOR_TESTS = "ci/ci_workspace_for_tests"
   static final GIT_FOLDER_FOR_TESTS = "git_folder_for_tests"
 
-  @Rule
-  public final EnvironmentVariables environmentVariables = new EnvironmentVariables()
-
-  @Rule
-  public final RestoreSystemProperties restoreSystemProperties = new RestoreSystemProperties()
+  protected ControllableEnvironmentVariables env = ControllableEnvironmentVariables.setup()
 
   protected final String localFSGitWorkspace = resolve(CI_WORKSPACE_PATH_FOR_TESTS)
 
-  def setup() {
-    // Clear all environment variables to avoid clashes between
-    // real CI/Git environment variables and the spec CI/Git
-    // environment variables.
-    environmentVariables.clear(System.getenv().keySet() as String[])
+  void cleanup() {
+    env.clear()
   }
 
   def "test ci provider info is set properly: #ciSpec.providerName #ciSpec.idx #ciSpec.testCaseName"() {
     setup:
     ciSpec.env.each {
-      environmentVariables.set(it.key, it.value.toString())
+      env.set(it.key, it.value.toString())
       if (it.key == "HOME") {
         System.setProperty("user.home", it.value)
       }
     }
 
     when:
-    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(System.getenv()))
+    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(env.getAll()))
     def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(getWorkspacePath())
     def ciInfo = ciProviderInfo.buildCIInfo()
     def prInfo = ciProviderInfo.buildPullRequestInfo()
@@ -68,13 +58,13 @@ abstract class CITagsProviderTest extends Specification {
   def "test user supplied commit hash takes precedence over auto-detected git info"() {
     setup:
     buildRemoteGitInfoEmpty().each {
-      environmentVariables.set(it.key, it.value)
+      env.set(it.key, it.value)
     }
 
-    environmentVariables.set(propertyNameToEnvironmentVariableName(UserSuppliedGitInfoBuilder.DD_GIT_COMMIT_SHA), "1234567890123456789012345678901234567890")
+    env.set(propertyNameToEnvironmentVariableName(UserSuppliedGitInfoBuilder.DD_GIT_COMMIT_SHA), "1234567890123456789012345678901234567890")
 
     when:
-    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(System.getenv()))
+    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(env.getAll()))
     def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(getWorkspacePath())
     def ciInfo = ciProviderInfo.buildCIInfo()
     def ciTagsProvider = ciTagsProvider()
@@ -87,13 +77,13 @@ abstract class CITagsProviderTest extends Specification {
   def "test user supplied repo url takes precedence over auto-detected git info"() {
     setup:
     buildRemoteGitInfoEmpty().each {
-      environmentVariables.set(it.key, it.value)
+      env.set(it.key, it.value)
     }
 
-    environmentVariables.set(propertyNameToEnvironmentVariableName(UserSuppliedGitInfoBuilder.DD_GIT_REPOSITORY_URL), "local supplied repo url")
+    env.set(propertyNameToEnvironmentVariableName(UserSuppliedGitInfoBuilder.DD_GIT_REPOSITORY_URL), "local supplied repo url")
 
     when:
-    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(System.getenv()))
+    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(env.getAll()))
     def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(getWorkspacePath())
     def ciInfo = ciProviderInfo.buildCIInfo()
     def ciTagsProvider = ciTagsProvider()
@@ -106,11 +96,11 @@ abstract class CITagsProviderTest extends Specification {
   def "test set local git info if remote git info is not present"() {
     setup:
     buildRemoteGitInfoEmpty().each {
-      environmentVariables.set(it.key, it.value)
+      env.set(it.key, it.value)
     }
 
     when:
-    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(System.getenv()))
+    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(env.getAll()))
     def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(getWorkspacePath())
     def ciInfo = ciProviderInfo.buildCIInfo()
     def ciTagsProvider = ciTagsProvider()
@@ -134,15 +124,15 @@ abstract class CITagsProviderTest extends Specification {
   def "test avoid setting local git info if remote commit does not match"() {
     setup:
     buildRemoteGitInfoMismatchLocalGit().each {
-      environmentVariables.set(it.key, it.value)
+      env.set(it.key, it.value)
     }
 
     when:
     def ciTagsProvider = ciTagsProvider()
 
     then:
     if (isCi()) {
-      CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(System.getenv()))
+      CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(env.getAll()))
       def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(getWorkspacePath())
       def ciInfo = ciProviderInfo.buildCIInfo()
       def tags = ciTagsProvider.getCiTags(ciInfo, PullRequestInfo.EMPTY)
@@ -176,7 +166,7 @@ abstract class CITagsProviderTest extends Specification {
 
     GitInfoProvider gitInfoProvider = new GitInfoProvider()
     gitInfoProvider.registerGitInfoBuilder(new UserSuppliedGitInfoBuilder())
-    gitInfoProvider.registerGitInfoBuilder(new CIProviderGitInfoBuilder(Config.get(), new CiEnvironmentImpl(System.getenv())))
+    gitInfoProvider.registerGitInfoBuilder(new CIProviderGitInfoBuilder(Config.get(), new CiEnvironmentImpl(env.getAll())))
     gitInfoProvider.registerGitInfoBuilder(new CILocalGitInfoBuilder(gitClientFactory, GIT_FOLDER_FOR_TESTS))
     return new CITagsProvider(gitInfoProvider)
   }

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/ci/GithubActionsInfoTest.groovy

@@ -33,11 +33,11 @@ class GithubActionsInfoTest extends CITagsProviderTest {
     def githubEvent = GithubActionsInfoTest.getResource("/ci/github-event.json")
     def githubEventPath = Paths.get(githubEvent.toURI())
 
-    environmentVariables.set(GithubActionsInfo.GITHUB_BASE_REF, "base-ref")
-    environmentVariables.set(GithubActionsInfo.GITHUB_EVENT_PATH, githubEventPath.toString())
+    env.set(GithubActionsInfo.GITHUB_BASE_REF, "base-ref")
+    env.set(GithubActionsInfo.GITHUB_EVENT_PATH, githubEventPath.toString())
 
     when:
-    def pullRequestInfo = new GithubActionsInfo(new CiEnvironmentImpl(System.getenv())).buildPullRequestInfo()
+    def pullRequestInfo = new GithubActionsInfo(new CiEnvironmentImpl(env.getAll())).buildPullRequestInfo()
 
     then:
     pullRequestInfo.getBaseBranch() == "base-ref"

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/ci/UnknownCIInfoTest.groovy

@@ -45,7 +45,7 @@ class UnknownCIInfoTest extends CITagsProviderTest {
     ]
 
     when:
-    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(System.getenv()))
+    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), GIT_FOLDER_FOR_TESTS, new CiEnvironmentImpl(env.getAll()))
     def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(workspaceForTests)
     def ciInfo = ciProviderInfo.buildCIInfo()
     def ciTagsProvider = ciTagsProvider()
@@ -62,9 +62,9 @@ class UnknownCIInfoTest extends CITagsProviderTest {
 
     GitInfoProvider gitInfoProvider = new GitInfoProvider()
     gitInfoProvider.registerGitInfoBuilder(new UserSuppliedGitInfoBuilder())
-    gitInfoProvider.registerGitInfoBuilder(new CIProviderGitInfoBuilder(Config.get(), new CiEnvironmentImpl(System.getenv())))
+    gitInfoProvider.registerGitInfoBuilder(new CIProviderGitInfoBuilder(Config.get(), new CiEnvironmentImpl(env.getAll())))
     gitInfoProvider.registerGitInfoBuilder(new CILocalGitInfoBuilder(gitClientFactory, "this-target-folder-does-not-exist"))
-    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), "this-target-folder-does-not-exist", new CiEnvironmentImpl(System.getenv()))
+    CIProviderInfoFactory ciProviderInfoFactory = new CIProviderInfoFactory(Config.get(), "this-target-folder-does-not-exist", new CiEnvironmentImpl(env.getAll()))
 
     def ciProviderInfo = ciProviderInfoFactory.createCIProviderInfo(workspaceForTests)
     def ciInfo = ciProviderInfo.buildCIInfo()

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/git/CIProviderGitInfoBuilderTest.groovy

@@ -2,25 +2,19 @@ package datadog.trace.civisibility.git
 
 import datadog.trace.api.Config
 import datadog.trace.civisibility.ci.env.CiEnvironmentImpl
-import org.junit.Rule
-import org.junit.contrib.java.lang.system.EnvironmentVariables
+import datadog.trace.test.util.ControllableEnvironmentVariables
 import spock.lang.Specification
 
 class CIProviderGitInfoBuilderTest extends Specification {
+  protected ControllableEnvironmentVariables env = ControllableEnvironmentVariables.setup()
 
-  @Rule
-  public final EnvironmentVariables environmentVariables = new EnvironmentVariables()
-
-  def setup() {
-    // Clear all environment variables to avoid clashes between
-    // real CI/Git environment variables and the spec CI/Git
-    // environment variables.
-    environmentVariables.clear(System.getenv().keySet() as String[])
+  void cleanup() {
+    env.clear()
   }
 
   def "test builds empty git info in an unknown repository"() {
     setup:
-    def builder = new CIProviderGitInfoBuilder(Config.get(), new CiEnvironmentImpl(System.getenv()))
+    def builder = new CIProviderGitInfoBuilder(Config.get(), new CiEnvironmentImpl(env.getAll()))
 
     when:
     def gitInfo = builder.build(null)

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerAgentTest.java

@@ -12,7 +12,6 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
-import static utils.TestHelper.setEnvVar;
 import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.util.RemoteConfigHelper;
@@ -22,6 +21,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.git.GitInfoProvider;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.test.util.ControllableEnvironmentVariables;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStreamReader;
@@ -54,6 +54,8 @@ public class DebuggerAgentTest {
   final MockWebServer server = new MockWebServer();
   HttpUrl url;
 
+  private ControllableEnvironmentVariables env = ControllableEnvironmentVariables.setup();
+
   private static void setFieldInContainerInfo(
       ContainerInfo containerInfo, String fieldName, Object value) {
     try {
@@ -67,6 +69,7 @@ private static void setFieldInContainerInfo(
 
   @BeforeEach
   public void setUp() {
+    env.clear();
     url = server.url(URL_PATH);
   }
 
@@ -175,14 +178,14 @@ public void tags() {
     when(config.getGlobalTags()).thenReturn(globalTags);
     // set env vars now to be cached by GitInfoProvider
     GitInfoProvider.INSTANCE.invalidateCache();
-    setEnvVar("DD_GIT_COMMIT_SHA", "sha1");
-    setEnvVar("DD_GIT_REPOSITORY_URL", "http://github.com");
+    env.set("DD_GIT_COMMIT_SHA", "sha1");
+    env.set("DD_GIT_REPOSITORY_URL", "http://github.com");
     String tags;
     try {
       tags = DebuggerAgent.getDefaultTagsMergedWithGlobalTags(config);
     } finally {
-      setEnvVar("DD_GIT_COMMIT_SHA", null);
-      setEnvVar("DD_GIT_REPOSITORY_URL", null);
+      env.set("DD_GIT_COMMIT_SHA", null);
+      env.set("DD_GIT_REPOSITORY_URL", null);
       GitInfoProvider.INSTANCE.invalidateCache();
     }
     Map<String, String> resultTags = new HashMap<>();

dd-java-agent/agent-debugger/src/test/java/utils/TestHelper.java

@@ -9,7 +9,6 @@
 import java.nio.file.Paths;
 import java.time.Duration;
 import java.util.List;
-import java.util.Map;
 import java.util.function.BooleanSupplier;
 
 public class TestHelper {
@@ -47,19 +46,4 @@ public static void assertWithTimeout(BooleanSupplier predicate, Duration timeout
     }
     assertTrue(predicate.getAsBoolean());
   }
-
-  public static void setEnvVar(String envName, String envValue) {
-    try {
-      Class<?> classOfMap = System.getenv().getClass();
-      Field field = classOfMap.getDeclaredField("m");
-      field.setAccessible(true);
-      if (envValue == null) {
-        ((Map<String, String>) field.get(System.getenv())).remove(envName);
-      } else {
-        ((Map<String, String>) field.get(System.getenv())).put(envName, envValue);
-      }
-    } catch (Exception ex) {
-      ex.printStackTrace();
-    }
-  }
 }