Adding utility to mock environment variables

apm-agent-core/pom.xml

@@ -140,13 +140,6 @@
             <version>1.2</version>
             <scope>test</scope>
         </dependency>
-        <!--For setting environment variables-->
-        <dependency>
-            <groupId>com.github.stefanbirkner</groupId>
-            <artifactId>system-lambda</artifactId>
-            <version>1.1.0</version>
-            <scope>test</scope>
-        </dependency>
         <dependency>
             <groupId>org.apache.httpcomponents</groupId>
             <artifactId>httpclient</artifactId>

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/context/MessageTest.java

@@ -1,3 +1,21 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.apm.agent.impl.context;
 
 import org.junit.jupiter.api.Test;

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/payload/ContainerInfoTest.java

@@ -18,16 +18,19 @@
  */
 package co.elastic.apm.agent.impl.payload;
 
+import co.elastic.apm.agent.util.CustomEnvVariables;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.DisabledOnJre;
 import org.junit.jupiter.api.condition.JRE;
 
 import javax.annotation.Nullable;
 
-import static com.github.stefanbirkner.systemlambda.SystemLambda.withEnvironmentVariable;
+import java.util.HashMap;
+import java.util.Map;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
-public class ContainerInfoTest {
+public class ContainerInfoTest extends CustomEnvVariables {
 
     @Test
     void testContainerIdParsing() {
@@ -142,25 +145,25 @@ void testKubernetesDownwardApi() throws Exception {
         String podName = "downward-api-pod-name";
         String nodeName = "downward-api-node-name";
         String namespace = "downward-api-namespace";
-        withEnvironmentVariable("KUBERNETES_NODE_NAME", nodeName)
-            .and("KUBERNETES_POD_NAME", podName)
-            .and("KUBERNETES_NAMESPACE", namespace)
-            .and("KUBERNETES_POD_UID", podUid)
-            .execute(systemInfo::findContainerDetails);
+
+        Map<String, String> mockedEnv = new HashMap<>();
+        mockedEnv.put("KUBERNETES_NODE_NAME", nodeName);
+        mockedEnv.put("KUBERNETES_POD_NAME", podName);
+        mockedEnv.put("KUBERNETES_NAMESPACE", namespace);
+        mockedEnv.put("KUBERNETES_POD_UID", podUid);
+        runWithCustomEnvVariables(mockedEnv, systemInfo::findContainerDetails);
         assertKubernetesInfo(systemInfo, podUid, podName, nodeName, namespace);
 
         // test partial settings
         systemInfo = assertContainerId(line, containerId);
         assertKubernetesInfo(systemInfo, originalPodUid, hostName, null, null);
-        withEnvironmentVariable("KUBERNETES_NODE_NAME", nodeName)
-            .and("KUBERNETES_POD_NAME", null)
-            .and("KUBERNETES_NAMESPACE", namespace)
-            .and("KUBERNETES_POD_UID", null)
-            .execute(systemInfo::findContainerDetails);
-        systemInfo.findContainerDetails();
+        mockedEnv.put("KUBERNETES_POD_NAME", null);
+        mockedEnv.put("KUBERNETES_POD_UID", null);
+        runWithCustomEnvVariables(mockedEnv, systemInfo::findContainerDetails);
         assertKubernetesInfo(systemInfo, originalPodUid, hostName, nodeName, namespace);
     }
 
+
     private SystemInfo createSystemInfo() {
         return new SystemInfo("arch", "my-host", "platform");
     }

apm-agent-core/src/test/java/co/elastic/apm/agent/testinstr/SystemAllEnvVariablesInstrumentation.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.apm.agent.testinstr;
+
+import co.elastic.apm.agent.sdk.advice.AssignTo;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
+
+public class SystemAllEnvVariablesInstrumentation extends SystemEnvVariableInstrumentation {
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return isStatic().and(named("getenv").and(takesNoArguments()));
+    }
+
+    public static class AdviceClass {
+        @AssignTo.Return
+        @Advice.OnMethodExit(onThrowable = Throwable.class, inline = false)
+        public static Map<String, String> appendToEnvVariables(@Advice.Return Map<String, String> ret) {
+            Map<String, String> customEnvVariables = customEnvVariablesTL.get();
+            if (customEnvVariables != null && !customEnvVariables.isEmpty()) {
+                ret = new HashMap<>(ret);
+                ret.putAll(customEnvVariables);
+            }
+            return ret;
+        }
+    }
+}

apm-agent-core/src/test/java/co/elastic/apm/agent/testinstr/SystemEnvVariableInstrumentation.java

@@ -0,0 +1,79 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.apm.agent.testinstr;
+
+import co.elastic.apm.agent.bci.TracerAwareInstrumentation;
+import co.elastic.apm.agent.sdk.advice.AssignTo;
+import co.elastic.apm.agent.sdk.state.GlobalVariables;
+import co.elastic.apm.agent.sdk.weakconcurrent.DetachedThreadLocal;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+public abstract class SystemEnvVariableInstrumentation extends TracerAwareInstrumentation {
+
+    protected static final DetachedThreadLocal<Map<String, String>> customEnvVariablesTL =
+        GlobalVariables.get(SystemEnvVariableInstrumentation.class, "customEnvVariables", WeakConcurrent.<Map<String, String>>buildThreadLocal());
+
+    /**
+     * Sets custom env variables that will be added to the actual env variables returned by {@link System#getenv()} or
+     * {@link System#getenv(String)} on the current thread.
+     * NOTE: caller must clear the custom variables when they are not required anymore through {@link #clearCustomEnvVariables()}.
+     * @param customEnvVariables a map of key-value pairs that will be appended to the actual environment variables
+     *                           returned by {@link System#getenv()} on the current thread
+     */
+    public static void setCustomEnvVariables(Map<String, String> customEnvVariables) {
+        customEnvVariablesTL.set(customEnvVariables);
+    }
+
+    public static void clearCustomEnvVariables() {
+        customEnvVariablesTL.remove();
+    }
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return named("java.lang.System");
+    }
+
+    @Override
+    public Collection<String> getInstrumentationGroupNames() {
+        return Collections.emptyList();
+    }
+
+    public static class AdviceClass {
+        @AssignTo.Return
+        @Advice.OnMethodExit(onThrowable = Throwable.class, inline = false)
+        public static Map<String, String> appendToEnvVariables(@Advice.Return Map<String, String> ret) {
+            Map<String, String> customEnvVariables = customEnvVariablesTL.get();
+            if (customEnvVariables != null && !customEnvVariables.isEmpty()) {
+                ret = new HashMap<>(ret);
+                ret.putAll(customEnvVariables);
+            }
+            return ret;
+        }
+    }
+}

apm-agent-core/src/test/java/co/elastic/apm/agent/testinstr/SystemSingleEnvVariablesInstrumentation.java

@@ -0,0 +1,53 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.apm.agent.testinstr;
+
+import co.elastic.apm.agent.sdk.advice.AssignTo;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import java.util.Map;
+
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+public class SystemSingleEnvVariablesInstrumentation extends SystemEnvVariableInstrumentation {
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return isStatic().and(named("getenv").and(takesArguments(1)));
+    }
+
+    public static class AdviceClass {
+        @AssignTo.Return
+        @Advice.OnMethodExit(onThrowable = Throwable.class, inline = false)
+        public static String appendToEnvVariables(@Advice.Argument(0) String varName, @Advice.Return String ret) {
+            Map<String, String> customEnvVariables = customEnvVariablesTL.get();
+            if (customEnvVariables != null) {
+                String customValue = customEnvVariables.get(varName);
+                if (customValue != null) {
+                    ret = customValue;
+                }
+            }
+            return ret;
+        }
+    }
+}

apm-agent-core/src/test/java/co/elastic/apm/agent/util/CustomEnvVariables.java

@@ -0,0 +1,67 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.apm.agent.util;
+
+import co.elastic.apm.agent.AbstractInstrumentationTest;
+import co.elastic.apm.agent.testinstr.SystemSingleEnvVariablesInstrumentation;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class CustomEnvVariables extends AbstractInstrumentationTest {
+
+    protected void runWithCustomEnvVariables(Map<String, String> customEnvVariables, Runnable task) {
+        try {
+            SystemSingleEnvVariablesInstrumentation.setCustomEnvVariables(customEnvVariables);
+            task.run();
+        } finally {
+            SystemSingleEnvVariablesInstrumentation.clearCustomEnvVariables();
+        }
+    }
+
+    @Test
+    void testCustomSingleEnvVariable() {
+        String pathVariable = "PATH";
+        final String originalPath = System.getenv(pathVariable);
+        String mockPath = "mock/path";
+        final Map<String, String> customVariables = Map.of("key1", "value1", pathVariable, mockPath);
+        runWithCustomEnvVariables(customVariables, () -> {
+            String returnedPath = System.getenv(pathVariable);
+            assertThat(returnedPath).isEqualTo(mockPath);
+        });
+        String returnedPath = System.getenv(pathVariable);
+        assertThat(returnedPath).isEqualTo(originalPath);
+    }
+
+    @Test
+    void testSingleEnvVariables() {
+        final Map<String, String> originalVariables = System.getenv();
+        final Map<String, String> customVariables = Map.of("key1", "value1", "key2", "value2");
+        runWithCustomEnvVariables(customVariables, () -> {
+            Map<String, String> returnedEnvVariables = System.getenv();
+            assertThat(returnedEnvVariables).containsAllEntriesOf(originalVariables);
+            assertThat(returnedEnvVariables).containsAllEntriesOf(customVariables);
+        });
+        Map<String, String> returnedEnvVariables = System.getenv();
+        assertThat(returnedEnvVariables).containsAllEntriesOf(originalVariables);
+        customVariables.forEach((key, value) -> assertThat(returnedEnvVariables).doesNotContainEntry(key, value));
+    }
+}

apm-agent-core/src/test/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation

@@ -0,0 +1,2 @@
+co.elastic.apm.agent.testinstr.SystemSingleEnvVariablesInstrumentation
+co.elastic.apm.agent.testinstr.SystemAllEnvVariablesInstrumentation