From 838dab0ae4f85cc70799b9fdfacfb9097628610d Mon Sep 17 00:00:00 2001
From: Pawel Sadlo <psadlo@virtuslab.com>
Date: Sat, 21 Sep 2024 09:07:16 +0200
Subject: [PATCH] fixing tests

---
 .../src/mill/initmodule/InitModule.scala      | 44 ++++++++++++-------
 .../feature/init/src/MillInitTests.scala      | 18 ++++++--
 main/src/mill/main/MainModule.scala           |  9 ++--
 3 files changed, 48 insertions(+), 23 deletions(-)

diff --git a/initmodule/src/mill/initmodule/InitModule.scala b/initmodule/src/mill/initmodule/InitModule.scala
index 2b64aa84e3e..853931e5aac 100644
--- a/initmodule/src/mill/initmodule/InitModule.scala
+++ b/initmodule/src/mill/initmodule/InitModule.scala
@@ -2,6 +2,7 @@ package mill.initmodule
 
 import mill.define.{Discover, ExternalModule}
 import mill.{Command, Module, T}
+import os.Path
 
 import java.io.IOException
 import scala.util.{Failure, Success, Try, Using}
@@ -11,36 +12,49 @@ object InitModule extends ExternalModule with InitModule {
 }
 
 trait InitModule extends Module {
-  def init(@mainargs.arg(positional = true, short = 'e') example: Option[String]): Command[String] =
-    T.command {
-      Using(getClass.getClassLoader.getResourceAsStream("exampleList.txt")) { exampleList =>
-        val reader = upickle.default.reader[Seq[(String, String)]]
-        val exampleNames: Seq[(String, String)] = upickle.default.read(exampleList)(reader)
 
-        val result: Try[String] = example match {
+  type ExampleUrl = String
+  type ExampleId = String
+
+  private def usingExamples[T](fun: Seq[(ExampleId, ExampleUrl)] => T): Try[T] =
+    Using(getClass.getClassLoader.getResourceAsStream("exampleList.txt")) { exampleList =>
+      val reader = upickle.default.reader[Seq[(ExampleId, ExampleUrl)]]
+      val exampleNames: Seq[(ExampleId, ExampleUrl)] = upickle.default.read(exampleList)(reader)
+      fun(exampleNames)
+    }
+
+  /**
+   * @return Seq of example names or Seq with path to parent dir where downloaded example was unpacked
+   */
+  def init(@mainargs.arg(positional = true, short = 'e') exampleId: Option[ExampleId])
+      : Command[Seq[String]] =
+    T.command {
+      usingExamples { examples =>
+        val result: Try[(Seq[String], String)] = exampleId match {
           case None =>
+            val exampleIds: Seq[ExampleId] = examples.map { case (exampleId, _) => exampleId }
             val msg =
               "Run init with one of the following examples as an argument to download and extract example:\n"
-            val examplesMsgPart = exampleNames.map { case (path, _) => path }.mkString("\n")
-            val out = msg + examplesMsgPart
-            Success(out)
+            val message = msg + exampleIds.mkString("\n")
+            Success(exampleIds, message)
           case Some(value) =>
-            for {
-              url <- exampleNames.toMap.get(value).toRight(new Exception(
+            val result: Try[(Seq[String], String)] = for {
+              url <- examples.toMap.get(value).toRight(new Exception(
                 s"Example [$value] is not present in examples list"
               )).toTry
               path <- Try(mill.util.Util.downloadUnpackZip(url, os.rel)(T.workspace)).recoverWith(
                 _ => Failure(new IOException(s"Couldn't download example: [$value]"))
               )
               _ = Try(os.remove(T.workspace / "tmp.zip"))
-            } yield "Example downloaded to " + path
+            } yield (Seq(path.path.toString()), s"Example downloaded to [${path.path.toString}]")
 
+            result
         }
         result
       }.flatten match {
-        case Success(value) =>
-          T.log.outputStream.println(value)
-          value
+        case Success((ret, msg)) =>
+          T.log.outputStream.println(msg)
+          ret
         case Failure(exception) =>
           T.log.error(exception.getMessage)
           throw exception
diff --git a/integration/feature/init/src/MillInitTests.scala b/integration/feature/init/src/MillInitTests.scala
index fcc30b3328a..b291e295965 100644
--- a/integration/feature/init/src/MillInitTests.scala
+++ b/integration/feature/init/src/MillInitTests.scala
@@ -10,12 +10,24 @@ object MillInitTests extends UtestIntegrationTestSuite {
       import tester._
       val res = eval("init")
       res.isSuccess ==> true
-      val firstProj = res.out.split("\n")(1)
+
+      val exampleListOut = out("init")
+      val parsed = exampleListOut.json.arr.map(_.str)
+
+      val firstProj = parsed.head
       val downloaded = eval(("init", firstProj))
       downloaded.isSuccess ==> true
-      val paths = os.walk(path = workspacePath, maxDepth = 1)
+
+      val outDir = out("init")
+      val parsedOutDir = outDir.json.arr.map(_.str).head
+
+      val downloadedParentDir = os.Path(parsedOutDir)
+
+      val paths = os.walk(path = downloadedParentDir, maxDepth = 1)
       val extractedDir = paths.find(_.last.endsWith(firstProj.replace("/", "-")))
-      extractedDir.isDefined ==> true
+
+      assert(extractedDir.isDefined)
+
       assert(os.exists(extractedDir.get / "build.mill"))
     }
   }
diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
index 19df503563d..81b8a013140 100644
--- a/main/src/mill/main/MainModule.scala
+++ b/main/src/mill/main/MainModule.scala
@@ -397,17 +397,16 @@ trait MainModule extends BaseModule0 {
   /**
    * The `init` command downloads specified example from mill releases page and extracts it to working directory.
    */
-  def init(evaluator: Evaluator, args: String*): Command[String] = Target.command {
+  def init(evaluator: Evaluator, args: String*): Command[ujson.Value] = Target.command {
     RunScript.evaluateTasksNamed(
       evaluator,
       Seq("mill.initmodule.InitModule/init") ++ args,
       SelectMode.Separated
     ) match {
-      case Right((_, Right(Seq((result, _))))) => result.toString
-      case Left(failStr) => failStr
-      case Right((_, Left(failStr))) => failStr
+      case Left(failStr) => throw new Exception(failStr)
+      case Right((_, Right(Seq((_, Some((_, jsonableResult))))))) => jsonableResult
+      case Right((_, Left(failStr))) => throw new Exception(failStr)
     }
-
   }
 
   private type VizWorker = (