Merge pull request #41 from webcompere/environment-other-mocks

Add further mocking to complete the range of things affected by the mock EnvironmentVariables

Author: ashleyfrieze

system-stubs-core/src/main/java/uk/org/webcompere/systemstubs/environment/EnvironmentVariableMocker.java

@@ -2,11 +2,13 @@
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
 
-import java.util.Map;
-import java.util.Stack;
+import java.util.*;
+import java.util.stream.Stream;
 
 import static java.util.stream.Collectors.toMap;
+import static java.util.stream.Collectors.toSet;
 
 /**
  * This takes control of the environment variables using {@link Mockito#mockStatic}. While there
@@ -18,14 +20,22 @@
 public class EnvironmentVariableMocker {
     private static final Stack<Map<String, String>> REPLACEMENT_ENV = new Stack<>();
 
+    private static final Set<String> MOCKED_METHODS = Stream.of("getenv", "environment", "toEnvironmentBlock")
+        .collect(toSet());
+
     static {
         try {
             Class<?> typeToMock = Class.forName("java.lang.ProcessEnvironment");
             Mockito.mockStatic(typeToMock, invocationOnMock -> {
-                if (REPLACEMENT_ENV.empty() || !invocationOnMock.getMethod().getName().equals("getenv")) {
+                if (REPLACEMENT_ENV.empty() || !MOCKED_METHODS.contains(invocationOnMock.getMethod().getName())) {
                     return invocationOnMock.callRealMethod();
                 }
-                Map<String, String> currentMockedEnvironment = REPLACEMENT_ENV.peek();
+
+                if ("toEnvironmentBlock".equals(invocationOnMock.getMethod().getName())) {
+                    return simulateToEnvironmentBlock(invocationOnMock);
+                }
+
+                Map<String, String> currentMockedEnvironment = getenv();
                 if (invocationOnMock.getMethod().getParameterCount() == 0) {
                     return filterNulls(currentMockedEnvironment);
                 }
@@ -38,6 +48,101 @@ public class EnvironmentVariableMocker {
         }
     }
 
+    /**
+     * The equivalent of <code>getenv</code> in the original ProcessEnvironment, assuming that
+     * mocking is "turned on"
+     * @return the current effective environment
+     */
+    private static Map<String, String> getenv() {
+        return REPLACEMENT_ENV.peek();
+    }
+
+    /**
+     * On Windows, this returns <code>String</code> and converts the inbound Map. On Linux/Mac
+     * this takes a second parameter and returns <code>byte[]</code>. Implementations ripped
+     * from the JDK source implementation
+     * @param invocationOnMock the call to the mocked <code>ProcessEnvironment</code>
+     * @return the environment serialized for the platform
+     */
+    private static Object simulateToEnvironmentBlock(InvocationOnMock invocationOnMock) {
+        if (invocationOnMock.getArguments().length == 1) {
+            return toEnvironmentBlockWindows(invocationOnMock.getArgument(0));
+        } else {
+            return toEnvironmentBlockNix(invocationOnMock.getArgument(0),
+                invocationOnMock.getArgument(1, int[].class));
+        }
+    }
+
+    /**
+     * Ripped from the JDK implementatioin
+     * @param m the map to convert
+     * @return string representation
+     */
+    private static String toEnvironmentBlockWindows(Map<String, String> m) {
+        // Sort Unicode-case-insensitively by name
+        List<Map.Entry<String,String>> list = new ArrayList<>(m.entrySet());
+        Collections.sort(list, (e1, e2) -> NameComparator.compareNames(e1.getKey(), e2.getKey()));
+
+        StringBuilder sb = new StringBuilder(m.size() * 30);
+        int cmp = -1;
+
+        // Some versions of MSVCRT.DLL require SystemRoot to be set.
+        // So, we make sure that it is always set, even if not provided
+        // by the caller.
+        final String systemRoot = "SystemRoot";
+
+        for (Map.Entry<String,String> e : list) {
+            String key = e.getKey();
+            String value = e.getValue();
+            if (cmp < 0 && (cmp = NameComparator.compareNames(key, systemRoot)) > 0) {
+                // Not set, so add it here
+                addToEnvIfSet(sb, systemRoot);
+            }
+            addToEnv(sb, key, value);
+        }
+        if (cmp < 0) {
+            // Got to end of list and still not found
+            addToEnvIfSet(sb, systemRoot);
+        }
+        if (sb.length() == 0) {
+            // Environment was empty and SystemRoot not set in parent
+            sb.append('\u0000');
+        }
+        // Block is double NUL terminated
+        sb.append('\u0000');
+        return sb.toString();
+    }
+
+    // code taken from the original in ProcessEnvironment
+    @SuppressFBWarnings({"PZLA_PREFER_ZERO_LENGTH_ARRAYS", "DM_DEFAULT_ENCODING"})
+    private static byte[] toEnvironmentBlockNix(Map<String, String> m, int[] envc) {
+        if (m == null) {
+            return null;
+        }
+        int count = m.size() * 2; // For added '=' and NUL
+        for (Map.Entry<String, String> entry : m.entrySet()) {
+            count += entry.getKey().getBytes().length;
+            count += entry.getValue().getBytes().length;
+        }
+
+        byte[] block = new byte[count];
+
+        int i = 0;
+        for (Map.Entry<String, String> entry : m.entrySet()) {
+            final byte[] key   = entry.getKey().getBytes();
+            final byte[] value = entry.getValue().getBytes();
+            System.arraycopy(key, 0, block, i, key.length);
+            i += key.length;
+            block[i++] = (byte) '=';
+            System.arraycopy(value, 0, block, i, value.length);
+            i += value.length + 1;
+            // No need to write NUL byte explicitly
+            //block[i++] = (byte) '\u0000';
+        }
+        envc[0] = m.size();
+        return block;
+    }
+
     private static Map<String, String> filterNulls(Map<String, String> currentMockedEnvironment) {
         return currentMockedEnvironment.entrySet()
             .stream()
@@ -80,4 +185,48 @@ public static boolean pop() {
     public static boolean remove(Map<String, String> theOneToPop) {
         return REPLACEMENT_ENV.remove(theOneToPop);
     }
+
+    @SuppressFBWarnings("SE_COMPARATOR_SHOULD_BE_SERIALIZABLE")
+    private static final class NameComparator
+        implements Comparator<String> {
+
+        public int compare(String s1, String s2) {
+            return compareNames(s1, s2);
+        }
+
+        public static int compareNames(String s1, String s2) {
+            // We can't use String.compareToIgnoreCase since it
+            // canonicalizes to lower case, while Windows
+            // canonicalizes to upper case!  For example, "_" should
+            // sort *after* "Z", not before.
+            int n1 = s1.length();
+            int n2 = s2.length();
+            int min = Math.min(n1, n2);
+            for (int i = 0; i < min; i++) {
+                char c1 = s1.charAt(i);
+                char c2 = s2.charAt(i);
+                if (c1 != c2) {
+                    c1 = Character.toUpperCase(c1);
+                    c2 = Character.toUpperCase(c2);
+                    if (c1 != c2) {
+                        // No overflow because of numeric promotion
+                        return c1 - c2;
+                    }
+                }
+            }
+            return n1 - n2;
+        }
+    }
+
+    // add the environment variable to the child, if it exists in parent
+    private static void addToEnvIfSet(StringBuilder sb, String name) {
+        String s = getenv().get(name);
+        if (s != null) {
+            addToEnv(sb, name, s);
+        }
+    }
+
+    private static void addToEnv(StringBuilder sb, String name, String val) {
+        sb.append(name).append('=').append(val).append('\u0000');
+    }
 }

system-stubs-core/src/test/java/uk/org/webcompere/systemstubs/environment/EnvironmentVariableMockerTest.java

@@ -2,11 +2,18 @@
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledOnOs;
 
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.condition.OS.*;
 
 class EnvironmentVariableMockerTest {
 
@@ -62,4 +69,70 @@ void whenAddNullThenItIsNotInGetEnvReturn() {
 
         assertThat(System.getenv()).doesNotContainKey("foo");
     }
+
+    @Test
+    void processBuilderEnvironmentIsAffectedByMockEnvironment() {
+        Map<String, String> newMap = new HashMap<>();
+        newMap.put("FOO", "bar");
+        EnvironmentVariableMocker.connect(newMap);
+        assertThat(new ProcessBuilder().environment()).containsEntry("FOO", "bar");
+    }
+
+    @EnabledOnOs({ MAC, LINUX })
+    @Test
+    void processBuilderEnvironmentWhenLaunchingNewApplicationIsAffected() throws Exception {
+        Map<String, String> newMap = new HashMap<>();
+        newMap.put("FOO", "bar");
+
+        EnvironmentVariableMocker.connect(newMap);
+
+        ProcessBuilder builder = new ProcessBuilder("/usr/bin/env");
+        builder.environment().put("BING", "bong");
+        String output = executeProcessAndGetOutput(builder);
+
+        assertThat(output).contains("FOO=bar");
+        assertThat(output).contains("BING=bong");
+    }
+
+    @EnabledOnOs({ MAC, LINUX })
+    @Test
+    void canLaunchWithDefaultEnvironmentAndNothingIsAdded() throws Exception {
+        Map<String, String> newMap = new HashMap<>();
+
+        EnvironmentVariableMocker.connect(newMap);
+
+        ProcessBuilder builder = new ProcessBuilder("/usr/bin/env");
+        String output = executeProcessAndGetOutput(builder);
+        assertThat(output).doesNotContain("FOO=bar");
+        assertThat(output).doesNotContain("BING=bong");
+    }
+
+    @EnabledOnOs(WINDOWS)
+    @Test
+    void windowsProcessBuilderEnvironmentCanReadEnvironmentFromMock() throws Exception {
+        Map<String, String> newMap = new HashMap<>();
+        newMap.put("FOO", "bar");
+
+        EnvironmentVariableMocker.connect(newMap);
+
+        ProcessBuilder builder = new ProcessBuilder(Arrays.asList("cmd.exe", "/c", "set"));
+        builder.environment().put("BING", "bong");
+        String output = executeProcessAndGetOutput(builder);
+
+        assertThat(output).contains("FOO=bar");
+        assertThat(output).contains("BING=bong");
+    }
+
+    private String executeProcessAndGetOutput(ProcessBuilder builder) throws IOException {
+        Process process = builder.start();
+
+        try (BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
+            return bufferedReader.lines().collect(Collectors.joining("\n"));
+        } catch (Throwable t) {
+            System.err.println(t.getMessage());
+            throw t;
+        } finally {
+            process.destroy();
+        }
+    }
 }

system-stubs-core/src/test/java/uk/org/webcompere/systemstubs/environment/EnvironmentVariablesTest.java

@@ -173,4 +173,20 @@ void setVariablesWithVarArgsSetterOddInputs() {
             .set("a"))
             .isInstanceOf(IllegalArgumentException.class);
     }
+
+    @Test
+    void environmentVariablesWhenAccessingOtherWays() throws Exception {
+        new EnvironmentVariables(singletonMap("FOO", "bar"))
+            .execute(() -> {
+                assertThat(System.getenv()).containsEntry("FOO", "bar");
+            });
+    }
+
+    @Test
+    void environmentVariablesInProcessBuilder() throws Exception {
+        new EnvironmentVariables(singletonMap("FOO", "bar"))
+            .execute(() -> {
+                assertThat(new ProcessBuilder().environment()).containsEntry("FOO", "bar");
+            });
+    }
 }