JAVA-5781 Add env var override (#1623)

JAVA-5781

Author: katcharov

driver-core/src/main/com/mongodb/internal/connection/ClientMetadataHelper.java

@@ -224,7 +224,7 @@ private enum Orchestrator {
         K8S("kubernetes") {
             @Override
             boolean isCurrentOrchestrator() {
-                return System.getenv("KUBERNETES_SERVICE_HOST") != null;
+                return FaasEnvironment.getEnv("KUBERNETES_SERVICE_HOST") != null;
             }
         },
         UNKNOWN(null);

driver-core/src/main/com/mongodb/internal/connection/FaasEnvironment.java

@@ -20,7 +20,9 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 enum FaasEnvironment {
     AWS_LAMBDA("aws.lambda"),
@@ -29,21 +31,23 @@ enum FaasEnvironment {
     VERCEL("vercel"),
     UNKNOWN(null);
 
+    static final Map<String, String> ENV_OVERRIDES_FOR_TESTING = new HashMap<>();
+
     static FaasEnvironment getFaasEnvironment() {
         List<FaasEnvironment> result = new ArrayList<>();
-        String awsExecutionEnv = System.getenv("AWS_EXECUTION_ENV");
+        String awsExecutionEnv = getEnv("AWS_EXECUTION_ENV");
 
-        if (System.getenv("VERCEL") != null) {
+        if (getEnv("VERCEL") != null) {
             result.add(FaasEnvironment.VERCEL);
         }
         if ((awsExecutionEnv != null && awsExecutionEnv.startsWith("AWS_Lambda_"))
-                || System.getenv("AWS_LAMBDA_RUNTIME_API") != null) {
+            || getEnv("AWS_LAMBDA_RUNTIME_API") != null) {
             result.add(FaasEnvironment.AWS_LAMBDA);
         }
-        if (System.getenv("FUNCTIONS_WORKER_RUNTIME") != null) {
+        if (getEnv("FUNCTIONS_WORKER_RUNTIME") != null) {
             result.add(FaasEnvironment.AZURE_FUNC);
         }
-        if (System.getenv("K_SERVICE") != null || System.getenv("FUNCTION_NAME") != null) {
+        if (getEnv("K_SERVICE") != null || getEnv("FUNCTION_NAME") != null) {
             result.add(FaasEnvironment.GCP_FUNC);
         }
         // vercel takes precedence over aws.lambda
@@ -56,6 +60,14 @@ static FaasEnvironment getFaasEnvironment() {
         return result.get(0);
     }
 
+    @Nullable
+    public static String getEnv(final String key) {
+        if (ENV_OVERRIDES_FOR_TESTING.containsKey(key)) {
+            return ENV_OVERRIDES_FOR_TESTING.get(key);
+        }
+        return System.getenv(key);
+    }
+
     @Nullable
     private final String name;
 
@@ -95,11 +107,11 @@ public Integer getMemoryMb() {
     public String getRegion() {
         switch (this) {
             case AWS_LAMBDA:
-                return System.getenv("AWS_REGION");
+                return getEnv("AWS_REGION");
             case GCP_FUNC:
-                return System.getenv("FUNCTION_REGION");
+                return getEnv("FUNCTION_REGION");
             case VERCEL:
-                return System.getenv("VERCEL_REGION");
+                return getEnv("VERCEL_REGION");
             default:
                 return null;
         }
@@ -108,7 +120,7 @@ public String getRegion() {
     @Nullable
     private static Integer getEnvInteger(final String name) {
         try {
-            String value = System.getenv(name);
+            String value = getEnv(name);
             return Integer.parseInt(value);
         } catch (NumberFormatException e) {
             return null;

driver-core/src/test/functional/com/mongodb/client/WithWrapper.java

@@ -16,13 +16,11 @@
 
 package com.mongodb.client;
 
+import com.mongodb.internal.connection.FaasEnvironmentAccessor;
 import com.mongodb.lang.Nullable;
 
-import java.lang.reflect.Field;
 import java.util.Map;
 
-import static java.lang.System.getenv;
-
 @FunctionalInterface
 public interface WithWrapper {
 
@@ -34,7 +32,7 @@ static WithWrapper withWrapper() {
 
     default WithWrapper withEnvironmentVariable(final String name, @Nullable final String value) {
         return runnable -> {
-            Map<String, String> innerMap = getEnvInnerMap();
+            Map<String, String> innerMap = FaasEnvironmentAccessor.getFaasEnvMap();
             String original = innerMap.get(name);
             if (value == null) {
                 innerMap.remove(name);
@@ -65,16 +63,4 @@ default WithWrapper withSystemProperty(final String name, final String value) {
         };
     }
 
-    static Map<String, String> getEnvInnerMap() {
-        try {
-            Map<String, String> env = getenv();
-            Field field = env.getClass().getDeclaredField("m");
-            field.setAccessible(true);
-            @SuppressWarnings("unchecked")
-            Map<String, String> result = (Map<String, String>) field.get(env);
-            return result;
-        } catch (IllegalAccessException | NoSuchFieldException e) {
-            throw new RuntimeException(e);
-        }
-    }
 }

driver-core/src/test/functional/com/mongodb/internal/connection/FaasEnvironmentAccessor.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * 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.mongodb.internal.connection;
+
+import java.util.Map;
+
+/**
+ * In the same package as FaasEnvironment, to access package-private
+ */
+public final class FaasEnvironmentAccessor {
+    private FaasEnvironmentAccessor() {
+    }
+
+    public static Map<String, String> getFaasEnvMap() {
+        return FaasEnvironment.ENV_OVERRIDES_FOR_TESTING;
+    }
+}