Fix testing as Java application (#3432)

Hello, I'm currently working on integrating Mill over the BSP plugin for
IntelliJ (https://github.com/JetBrains/hirschgarten - now the BSP plugin
code is merged with Bazel-related plugins).

I have stumbled upon an issue, where the `Test <module.name> as Java
application` action did not work with a testing suite. It only performed
a Query JVM environment task and tests were not run (later on it turned
out that it was caused by the `zincWorker` in
`MillJvmBuildServer#jvmRunTestEnvironment`, which could not find the
correct test main class).

After digging into firstly the BSP plugin's code and later on Mill's
code, I discovered that in the
`MillJvmBuildServer#jvmRunTestEnvironment` there is testing logic
lacking. It did not manage to fetch the correct test-specific main
class, classpath and main arguments necessary for the testing as a local
Java application.

I figured out that in order to make it work, I needed to leverage the
`TestModule`'s logic - I reproduced the `MillJvmBuildServer#testTask`
logic of creating the java process' arguments and used it in a new
`TestModule#getTestEnvironmentVars` function, which is then used in the
`MillJvmBuildServer#jvmRunTestEnvironment` in case of a `TestModule with
JavaModule` module (note: the code reproduction I've done might not be
the most graceful solution - please let me know if I can do it better).
Then the test-specific main class, classpath and main arguments I
mentioned are assigned to adequate fields of the function's result -
`JvmEnvironmentItem`, which is finally used on the BSP plugin side,
which launches the actual local Java application.

I tested the solution with an example Mill project (generated with the
g8 template). With BSP plugin installed, I imported the project with the
new importing via BSP plugin option (it's still WIP and not yet
available in official Scala plugin distribution). Then I checked that
indeed the `Test <module.name> as Java application` gutter triggers the
tests and their result is displayed in the terminal.

If you need any help with e.g. setting up the BSP plugin, I'm happy to
help

bsp/worker/src/mill/bsp/worker/MillJvmBuildServer.scala

@@ -15,7 +15,8 @@ import ch.epfl.scala.bsp4j.{
 }
 import mill.T
 import mill.bsp.worker.Utils.sanitizeUri
-import mill.scalalib.JavaModule
+import mill.scalalib.api.CompilationResult
+import mill.scalalib.{JavaModule, TestModule}
 
 import java.util.concurrent.CompletableFuture
 import scala.jdk.CollectionConverters._
@@ -50,6 +51,10 @@ private trait MillJvmBuildServer extends JvmBuildServer { this: MillBuildServer
       targetIds = _ => targetIds,
       tasks = {
         case m: JavaModule =>
+          val moduleSpecificTask = m match {
+            case m: TestModule => m.getTestEnvironmentVars()
+            case _ => m.compile
+          }
           T.task {
             (
               m.runClasspath(),
@@ -58,18 +63,51 @@ private trait MillJvmBuildServer extends JvmBuildServer { this: MillBuildServer
               m.forkEnv(),
               m.mainClass(),
               m.zincWorker().worker(),
-              m.compile()
+              moduleSpecificTask()
             )
           }
       }
     ) {
-      // We ignore all non-JavaModule
+      case (
+            ev,
+            state,
+            id,
+            _: TestModule with JavaModule,
+            (
+              _,
+              forkArgs,
+              forkWorkingDir,
+              forkEnv,
+              _,
+              _,
+              testEnvVars: (String, String, String, Seq[String])
+            )
+          ) =>
+        val (mainClass, testRunnerClassPath, argsFile, classpath) = testEnvVars
+        val fullMainArgs: List[String] = List(testRunnerClassPath, argsFile)
+        val item = new JvmEnvironmentItem(
+          id,
+          classpath.asJava,
+          forkArgs.asJava,
+          forkWorkingDir.toString(),
+          forkEnv.asJava
+        )
+        item.setMainClasses(List(mainClass).map(new JvmMainClass(_, fullMainArgs.asJava)).asJava)
+        item
       case (
             ev,
             state,
             id,
             _: JavaModule,
-            (runClasspath, forkArgs, forkWorkingDir, forkEnv, mainClass, zincWorker, compile)
+            (
+              runClasspath,
+              forkArgs,
+              forkWorkingDir,
+              forkEnv,
+              mainClass,
+              zincWorker,
+              compile: CompilationResult
+            )
           ) =>
         val classpath = runClasspath.map(_.path).map(sanitizeUri)
         val item = new JvmEnvironmentItem(

scalalib/src/mill/scalalib/TestModule.scala

@@ -76,6 +76,12 @@ trait TestModule
       testTask(T.task { args }, T.task { Seq.empty[String] })()
     }
 
+  def getTestEnvironmentVars(args: String*): Command[(String, String, String, Seq[String])] = {
+    T.command {
+      getTestEnvironmentVarsTask(T.task { args })()
+    }
+  }
+
   /**
    * Args to be used by [[testCached]].
    */
@@ -124,6 +130,40 @@ trait TestModule
    */
   def testReportXml: T[Option[String]] = T(Some("test-report.xml"))
 
+  /**
+   * Returns a Tuple where the first element is the main-class, second and third are main-class-arguments and the forth is classpath
+   */
+  private def getTestEnvironmentVarsTask(args: Task[Seq[String]])
+      : Task[(String, String, String, Seq[String])] =
+    T.task {
+      val mainClass = "mill.testrunner.entrypoint.TestRunnerMain"
+      val outputPath = T.dest / "out.json"
+      val selectors = Seq.empty
+
+      val testArgs = TestArgs(
+        framework = testFramework(),
+        classpath = runClasspath().map(_.path),
+        arguments = args(),
+        sysProps = Map.empty,
+        outputPath = outputPath,
+        colored = T.log.colored,
+        testCp = testClasspath().map(_.path),
+        home = T.home,
+        globSelectors = selectors
+      )
+
+      val argsFile = T.dest / "testargs"
+      os.write(argsFile, upickle.default.write(testArgs))
+
+      val testRunnerClasspathArg =
+        zincWorker().scalalibClasspath()
+          .map(_.path.toNIO.toUri.toURL).mkString(",")
+
+      val cp = (runClasspath() ++ zincWorker().testrunnerEntrypointClasspath()).map(_.path.toString)
+
+      Result.Success((mainClass, testRunnerClasspathArg, argsFile.toString, cp))
+    }
+
   /**
    * Whether or not to use the test task destination folder as the working directory
    * when running tests. `true` means test subprocess run in the `.dest/sandbox` folder of