Adds test for client -> server map propagation

* Moved the `System.getenv` side effect to the end of the world
* Adds a test to make sure that the `Map[String, String]` gets
propagated correctly
* Adds a `Ctx.Env` trait for consistency

https://github.com/lihaoyi/mill/issues/257

5 files changed, 135 insertions, 44 deletions

diff --git a/clientserver/src/mill/clientserver/Client.java b/clientserver/src/mill/clientserver/Client.java
index c570a65c..ccabc24d 100644
--- a/clientserver/src/mill/clientserver/Client.java
+++ b/clientserver/src/mill/clientserver/Client.java
@@ -7,10 +7,7 @@ import java.net.Socket;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.channels.FileChannel;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Iterator;
-import java.util.Properties;
+import java.util.*;
 
 public class Client {
     static void initServer(String lockBase, boolean setJnaNoSys) throws IOException,URISyntaxException{
@@ -51,6 +48,7 @@ public class Client {
     }
     public static void main(String[] args) throws Exception{
         boolean setJnaNoSys = System.getProperty("jna.nosys") == null;
+        Map<String, String> env = System.getenv();
         if (setJnaNoSys) {
             System.setProperty("jna.nosys", "true");
         }
@@ -82,7 +80,8 @@ public class Client {
                         System.in,
                         System.out,
                         System.err,
-                        args
+                        args,
+                        env
                 );
                 System.exit(exitCode);
             }
@@ -97,11 +96,12 @@ public class Client {
                           InputStream stdin,
                           OutputStream stdout,
                           OutputStream stderr,
-                          String[] args) throws Exception{
+                          String[] args,
+                          Map<String, String> env) throws Exception{
 
         FileOutputStream f = new FileOutputStream(lockBase + "/run");
         ClientServer.writeArgs(System.console() != null, args, f);
-        ClientServer.writeEnv(f);
+        ClientServer.writeMap(env, f);
         f.close();
 
         boolean serverInit = false;
diff --git a/clientserver/src/mill/clientserver/ClientServer.java b/clientserver/src/mill/clientserver/ClientServer.java
index 45c99ea0..15c20f41 100644
--- a/clientserver/src/mill/clientserver/ClientServer.java
+++ b/clientserver/src/mill/clientserver/ClientServer.java
@@ -43,16 +43,15 @@ public class ClientServer {
      * server (as the server remains alive over the course of several runs and
      * does not see the environment changes the client would)
      */
-    public static void writeEnv(OutputStream argStream) throws IOException {
-        Map<String, String> env = System.getenv();
-        argStream.write(env.size());
-        for (Map.Entry<String, String> kv : env.entrySet()) {
+    public static void writeMap(Map<String, String> map, OutputStream argStream) throws IOException {
+        argStream.write(map.size());
+        for (Map.Entry<String, String> kv : map.entrySet()) {
             writeString(argStream, kv.getKey());
             writeString(argStream, kv.getValue());
         }
     }
 
-    public static Map<String, String> parseEnv(InputStream argStream) throws IOException {
+    public static Map<String, String> parseMap(InputStream argStream) throws IOException {
         Map<String, String> env = new HashMap<>();
         int mapLength = argStream.read();
         for (int i = 0; i < mapLength; i++) {
diff --git a/clientserver/src/mill/clientserver/Server.scala b/clientserver/src/mill/clientserver/Server.scala
index d9e67489..4deac55f 100644
--- a/clientserver/src/mill/clientserver/Server.scala
+++ b/clientserver/src/mill/clientserver/Server.scala
@@ -78,7 +78,7 @@ class Server[T](lockBase: String,
     val argStream = new FileInputStream(lockBase + "/run")
     val interactive = argStream.read() != 0;
     val args = ClientServer.parseArgs(argStream)
-    val env = ClientServer.parseEnv(argStream)
+    val env = ClientServer.parseMap(argStream)
     argStream.close()
 
     var done = false
diff --git a/clientserver/test/src/mill/clientserver/ClientServerTests.scala b/clientserver/test/src/mill/clientserver/ClientServerTests.scala
index 0b55d4bb..ac2063ef 100644
--- a/clientserver/test/src/mill/clientserver/ClientServerTests.scala
+++ b/clientserver/test/src/mill/clientserver/ClientServerTests.scala
@@ -2,6 +2,7 @@ package mill.clientserver
 import java.io._
 import java.nio.file.Path
 
+import scala.collection.JavaConverters._
 import utest._
 class EchoServer extends ServerMain[Int]{
   def main0(args: Array[String],
@@ -14,9 +15,16 @@ class EchoServer extends ServerMain[Int]{
 
     val reader = new BufferedReader(new InputStreamReader(stdin))
     val str = reader.readLine()
-    stdout.println(str + args(0))
+    if (args.nonEmpty){
+      stdout.println(str + args(0))
+    }
+    env.toSeq.sortBy(_._1).foreach{
+      case (key, value) => stdout.println(s"$key=$value")
+    }
     stdout.flush()
-    stderr.println(str.toUpperCase + args(0))
+    if (args.nonEmpty){
+      stderr.println(str.toUpperCase + args(0))
+    }
     stderr.flush()
     (true, None)
   }
@@ -36,37 +44,39 @@ object ClientServerTests extends TestSuite{
     (tmpDir, locks)
   }
 
+  def spawnEchoServer(tmpDir : Path, locks: Locks): Unit = {
+    new Thread(() => new Server(
+      tmpDir.toString,
+      new EchoServer(),
+      () => (),
+      1000,
+      locks
+    ).run()).start()
+  }
+
+  def runClientAux(tmpDir : Path, locks: Locks)
+                  (env : Map[String, String], args: Array[String]) = {
+    val (in, out, err) = initStreams()
+    Server.lockBlock(locks.clientLock){
+      Client.run(
+        tmpDir.toString,
+        () => spawnEchoServer(tmpDir, locks),
+        locks,
+        in,
+        out,
+        err,
+        args,
+        env.asJava
+      )
+      Thread.sleep(100)
+      (new String(out.toByteArray), new String(err.toByteArray))
+    }
+  }
+
   def tests = Tests{
     'hello - {
       val (tmpDir, locks) = init()
-
-      def spawnEchoServer(): Unit = {
-        new Thread(() => new Server(
-          tmpDir.toString,
-          new EchoServer(),
-          () => (),
-          1000,
-          locks
-        ).run()).start()
-      }
-
-
-      def runClient(arg: String) = {
-        val (in, out, err) = initStreams()
-        Server.lockBlock(locks.clientLock){
-          Client.run(
-            tmpDir.toString,
-            () => spawnEchoServer(),
-            locks,
-            in,
-            out,
-            err,
-            Array(arg)
-          )
-          Thread.sleep(100)
-          (new String(out.toByteArray), new String(err.toByteArray))
-        }
-      }
+      def runClient(s: String) = runClientAux(tmpDir, locks)(Map.empty, Array(s))
 
       // Make sure the simple "have the client start a server and
       // exchange one message" workflow works from end to end.
@@ -117,5 +127,83 @@ object ClientServerTests extends TestSuite{
         err3 == "HELLO World\n"
       )
     }
+
+    'envVars - {
+      val (tmpDir, locks) = init()
+
+      def runClient(env : Map[String, String]) = runClientAux(tmpDir, locks)(env, Array())
+
+      // Make sure the simple "have the client start a server and
+      // exchange one message" workflow works from end to end.
+
+      assert(
+        locks.clientLock.probe(),
+        locks.serverLock.probe(),
+        locks.processLock.probe()
+      )
+
+      def longString(s : String) = Array.fill(1000)(s).mkString
+      val b1000 = longString("b")
+      val c1000 = longString("c")
+      val a1000 = longString("a")
+
+      val env = Map(
+        "a" -> a1000,
+        "b" -> b1000,
+        "c" -> c1000
+      )
+
+
+      val (out1, err1) = runClient(env)
+      val expected = s"a=$a1000\nb=$b1000\nc=$c1000\n"
+
+      assert(
+        out1 == expected,
+        err1 == ""
+      )
+
+      // Give a bit of time for the server to release the lock and
+      // re-acquire it to signal to the client that it's done
+      Thread.sleep(100)
+
+      assert(
+        locks.clientLock.probe(),
+        !locks.serverLock.probe(),
+        !locks.processLock.probe()
+      )
+
+      val path = List(
+        "/Users/foo/Library/Haskell/bin",
+        "/usr/local/git/bin",
+        "/sw/bin/",
+        "/usr/local/bin",
+        "/usr/local/",
+        "/usr/local/sbin",
+        "/usr/local/mysql/bin",
+        "/usr/local/bin",
+        "/usr/bin",
+        "/bin",
+        "/usr/sbin",
+        "/sbin",
+        "/opt/X11/bin",
+        "/usr/local/MacGPG2/bin",
+        "/Library/TeX/texbin",
+        "/usr/local/bin/",
+        "/Users/foo/bin",
+        "/Users/foo/go/bin",
+        "~/.bloop"
+      )
+
+      val pathEnvVar = path.mkString(":")
+      val (out2, err2) = runClient(Map("PATH" -> pathEnvVar))
+
+      val expected2 = s"PATH=$pathEnvVar\n"
+
+      assert(
+        out2 == expected2,
+        err2 == ""
+      )
+
+    }
   }
 }
diff --git a/core/src/mill/util/Ctx.scala b/core/src/mill/util/Ctx.scala
index 88a8baec..6c8b2afb 100644
--- a/core/src/mill/util/Ctx.scala
+++ b/core/src/mill/util/Ctx.scala
@@ -23,6 +23,9 @@ object Ctx{
   trait Home{
     def home: Path
   }
+  trait Env{
+    def env: Map[String, String]
+  }
   object Log{
     implicit def logToCtx(l: Logger): Log = new Log { def log = l }
   }
@@ -41,7 +44,8 @@ class Ctx(val args: IndexedSeq[_],
   extends Ctx.Dest
   with Ctx.Log
   with Ctx.Args
-  with Ctx.Home{
+  with Ctx.Home
+  with Ctx.Env {
 
   def dest = dest0()
   def length = args.length