From f6ed921c0358149e73e1dc056ff42c296831a065 Mon Sep 17 00:00:00 2001
From: Alexandre Archambault <alexandre.archambault@gmail.com>
Date: Thu, 15 Feb 2024 19:44:28 +0100
Subject: [PATCH] Run tests via bootstraps on Windows

Rather than with manifest JARs, that confuse Ammonite / Almond
---
 project/settings.sc | 118 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 118 insertions(+)

diff --git a/project/settings.sc b/project/settings.sc
index ad16f844e..1e344144e 100644
--- a/project/settings.sc
+++ b/project/settings.sc
@@ -261,6 +261,124 @@ trait AlmondTestModule
     with AlmondScalacOptions {
   // with AlmondScala2Or3Module {
 
+  // originally based on https://github.com/com-lihaoyi/mill/blob/3335d2a2f7f33766a680e30df6a7d0dc0fbe08b3/scalalib/src/mill/scalalib/TestModule.scala#L80-L146
+  // goal here is to use a coursier bootstrap rather than a manifest JAR when testUseArgsFile is true
+  protected def testTask(
+    args: Task[Seq[String]],
+    globSelectors: Task[Seq[String]]
+  ): Task[(String, Seq[mill.testrunner.TestResult])] =
+    T.task {
+      val outputPath  = T.dest / "out.json"
+      val useArgsFile = testUseArgsFile()
+
+      val (jvmArgs, props: Map[String, String]) =
+        if (useArgsFile) {
+          val (props, jvmArgs) = forkArgs().partition(_.startsWith("-D"))
+          val sysProps =
+            props
+              .map(_.drop(2).split("[=]", 2))
+              .map {
+                case Array(k, v) => k -> v
+                case Array(k)    => k -> ""
+              }
+              .toMap
+
+          jvmArgs -> sysProps
+        }
+        else
+          forkArgs() -> Map()
+
+      val testArgs = mill.testrunner.TestArgs(
+        framework = testFramework(),
+        classpath = runClasspath().map(_.path),
+        arguments = args(),
+        sysProps = props,
+        outputPath = outputPath,
+        colored = T.log.colored,
+        testCp = compile().classes.path,
+        home = T.home,
+        globSelectors = globSelectors()
+      )
+
+      val testRunnerClasspathArg = zincWorker().scalalibClasspath()
+        .map(_.path.toNIO.toUri.toURL)
+        .mkString(",")
+
+      val argsFile = T.dest / "testargs"
+      os.write(argsFile, upickle.default.write(testArgs))
+      val mainArgs = Seq(testRunnerClasspathArg, argsFile.toString)
+
+      runSubprocess(
+        mainClass = "mill.testrunner.entrypoint.TestRunnerMain",
+        classPath = (runClasspath() ++ zincWorker().testrunnerEntrypointClasspath()).map(_.path),
+        jvmArgs = jvmArgs,
+        envArgs = forkEnv(),
+        mainArgs = mainArgs,
+        workingDir = forkWorkingDir(),
+        useCpPassingJar = useArgsFile
+      )
+
+      if (!os.exists(outputPath)) mill.api.Result.Failure("Test execution failed.")
+      else
+        try {
+          val jsonOutput = ujson.read(outputPath.toIO)
+          val (doneMsg, results) =
+            upickle.default.read[(String, Seq[mill.testrunner.TestResult])](jsonOutput)
+          TestModule.handleResults(doneMsg, results, Some(T.ctx()))
+        }
+        catch {
+          case e: Throwable =>
+            mill.api.Result.Failure("Test reporting failed: " + e)
+        }
+    }
+
+  private def createBootstrapJar(
+    dest: os.Path,
+    classPath: Agg[os.Path],
+    mainClass: String
+  ): Unit = {
+    import coursier.launcher._
+    val content = Seq(ClassLoaderContent.fromUrls(classPath.toSeq.map(_.toNIO.toUri.toASCIIString)))
+    val params  = Parameters.Bootstrap(content, mainClass)
+    BootstrapGenerator.generate(params, dest.toNIO)
+  }
+
+  // originally based on https://github.com/com-lihaoyi/mill/blob/3335d2a2f7f33766a680e30df6a7d0dc0fbe08b3/main/util/src/mill/util/Jvm.scala#L117-L154
+  // goal is to use coursier bootstraps rather than manifest JARs
+  // the former load JARs via standard URLClassLoader-s, which is better dealt with in Ammonite and almond
+  private def runSubprocess(
+    mainClass: String,
+    classPath: Agg[os.Path],
+    jvmArgs: Seq[String] = Seq.empty,
+    envArgs: Map[String, String] = Map.empty,
+    mainArgs: Seq[String] = Seq.empty,
+    workingDir: os.Path = null,
+    useCpPassingJar: Boolean = false
+  )(implicit ctx: mill.api.Ctx): Unit = {
+
+    val (cp, mainClass0) =
+      if (useCpPassingJar && !classPath.iterator.isEmpty) {
+        val passingJar = os.temp(prefix = "run-", suffix = ".jar", deleteOnExit = false)
+        ctx.log.debug(
+          s"Creating classpath passing jar '$passingJar' with Class-Path: ${classPath.iterator.map(_.toNIO.toUri.toURL.toExternalForm).mkString(" ")}"
+        )
+        createBootstrapJar(passingJar, classPath, mainClass)
+        (Agg(passingJar), "coursier.bootstrap.launcher.Launcher")
+      }
+      else
+        (classPath, mainClass)
+
+    val args =
+      Vector(mill.util.Jvm.javaExe) ++
+        jvmArgs ++
+        Vector("-cp", cp.iterator.mkString(java.io.File.pathSeparator), mainClass0) ++
+        mainArgs
+
+    ctx.log.debug(s"Run subprocess with args: ${args.map(a => s"'$a'").mkString(" ")}")
+
+    mill.util.Jvm.runSubprocess(args, envArgs, workingDir)
+  }
+
   def ivyDeps       = Agg(Deps.utest)
   def testFramework = "utest.runner.Framework"