Add unit test for main method discovery

9 files changed, 78 insertions, 23 deletions

diff --git a/docs/pages/1 - Intro to Mill.md b/docs/pages/1 - Intro to Mill.md
index 74c12a2e..8d2c4847 100644
--- a/docs/pages/1 - Intro to Mill.md
+++ b/docs/pages/1 - Intro to Mill.md
@@ -59,7 +59,7 @@ This can be run from the Bash shell via:
 ```bash
 $ mill foo.compile                 # compile sources into classfiles
 
-$ mill foo.runMain foo.ExampleMain # run the specified main method
+$ mill foo.run                     # run the main method, if any
 
 $ mill foo.jar                     # bundle the classfiles into a jar
 
diff --git a/scalalib/src/mill/scalalib/ScalaModule.scala b/scalalib/src/mill/scalalib/ScalaModule.scala
index 05c34249..d62cc300 100644
--- a/scalalib/src/mill/scalalib/ScalaModule.scala
+++ b/scalalib/src/mill/scalalib/ScalaModule.scala
@@ -24,19 +24,24 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
   }
   def scalaVersion: T[String]
 
-  def mainClass: T[Option[String]] = T{
-    discoverMainClasses() match {
-      case Seq(main) => Some(main)
-      case _         => None
+  def mainClass: T[Option[String]] = None
+
+  def finalMainClass: T[String] = T{
+    mainClass() match {
+      case Some(main) => Result.Success(main)
+      case None =>
+        mill.scalalib.ScalaWorkerApi.scalaWorker().discoverMainClasses(compile()) match {
+          case Seq() => Result.Failure("No main class specified or found")
+          case Seq(main) => Result.Success(main)
+          case mains =>
+            Result.Failure(
+              s"Multiple main classes found (${mains.mkString(",")}) " +
+              "please explicitly specify which one to use by overriding mainClass"
+            )
+        }
     }
   }
 
-  def discoverMainClasses: T[Seq[String]] = T{
-    Task.traverse(transitiveModuleDeps){ module => T.task {
-      mill.scalalib.ScalaWorkerApi.scalaWorker().discoverMainClasses(module.compile())
-    }}().flatten.distinct
-  }
-
   def ivyDeps = T{ Agg.empty[Dep] }
   def compileIvyDeps = T{ Agg.empty[Dep] }
   def scalacPluginIvyDeps = T{ Agg.empty[Dep] }
@@ -201,7 +206,7 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
 
   def runLocal(args: String*) = T.command {
     Jvm.runLocal(
-      mainClass().getOrElse(throw new RuntimeException("No mainClass provided!")),
+      finalMainClass(),
       runClasspath().map(_.path),
       args
     )
@@ -209,7 +214,7 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
 
   def run(args: String*) = T.command{
     Jvm.interactiveSubprocess(
-      mainClass().getOrElse(throw new RuntimeException("No mainClass provided!")),
+      finalMainClass(),
       runClasspath().map(_.path),
       forkArgs(),
       forkEnv(),
diff --git a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
index ec2489af..5b7aaa04 100644
--- a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
+++ b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
@@ -63,5 +63,6 @@ trait ScalaWorkerApi {
                args: Seq[String])
               (implicit ctx: mill.util.Ctx.Log): (String, Seq[Result])
 
-  def discoverMainClasses(compilationResult: CompilationResult)(implicit ctx: mill.util.Ctx): Seq[String]
+  def discoverMainClasses(compilationResult: CompilationResult)
+                         (implicit ctx: mill.util.Ctx): Seq[String]
 }
diff --git a/scalalib/test/resources/hello-world-no-main/core/src-2.10/Shim.scala b/scalalib/test/resources/hello-world-no-main/core/src-2.10/Shim.scala
new file mode 100644
index 00000000..025cff5c
--- /dev/null
+++ b/scalalib/test/resources/hello-world-no-main/core/src-2.10/Shim.scala
@@ -0,0 +1,5 @@
+object Shim{
+  def main(args: Array[String]): Unit = {
+    Main0(args(0), scala.util.Properties.versionNumberString + " rox")
+  }
+}
\ No newline at end of file
diff --git a/scalalib/test/resources/hello-world-no-main/core/src-2.11/Shim.scala b/scalalib/test/resources/hello-world-no-main/core/src-2.11/Shim.scala
new file mode 100644
index 00000000..d98a6de1
--- /dev/null
+++ b/scalalib/test/resources/hello-world-no-main/core/src-2.11/Shim.scala
@@ -0,0 +1,5 @@
+object Shim{
+  def main(args: Array[String]): Unit = {
+    Main0(args(0), scala.util.Properties.versionNumberString + " pwns")
+  }
+}
\ No newline at end of file
diff --git a/scalalib/test/resources/hello-world-no-main/core/src-2.12/Shim.scala b/scalalib/test/resources/hello-world-no-main/core/src-2.12/Shim.scala
new file mode 100644
index 00000000..1f3e1c0f
--- /dev/null
+++ b/scalalib/test/resources/hello-world-no-main/core/src-2.12/Shim.scala
@@ -0,0 +1,5 @@
+object Shim{
+  def main(args: Array[String]): Unit = {
+    Main0(args(0), scala.util.Properties.versionNumberString + " leet")
+  }
+}
\ No newline at end of file
diff --git a/scalalib/test/resources/hello-world-no-main/core/src/Main.scala b/scalalib/test/resources/hello-world-no-main/core/src/Main.scala
new file mode 100644
index 00000000..5f68fd3f
--- /dev/null
+++ b/scalalib/test/resources/hello-world-no-main/core/src/Main.scala
@@ -0,0 +1,11 @@
+import scala.collection._
+import java.nio.file.{Files, Paths}
+
+import Main.{args, greeting}
+object Main0{
+  def apply(s: String, greeting: String) = {
+    val resultPath = Paths.get(s)
+    Files.createDirectories(resultPath.getParent)
+    Files.write(resultPath, greeting.getBytes)
+  }
+}
diff --git a/scalalib/test/resources/hello-world-no-main/core/src/Result.scala b/scalalib/test/resources/hello-world-no-main/core/src/Result.scala
new file mode 100644
index 00000000..d7d29a51
--- /dev/null
+++ b/scalalib/test/resources/hello-world-no-main/core/src/Result.scala
@@ -0,0 +1,7 @@
+object Person {
+  def fromString(s: String): Person = {
+    val Array(name, age) = s.split(":")
+    Person(name, age.toInt)
+  }
+}
+case class Person(name: String, age: Int)
diff --git a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
index 945bec27..7a788b1f 100644
--- a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
+++ b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
@@ -122,7 +122,7 @@ object HelloWorldTests extends TestSuite {
     "Person$.class"
   )
 
-  def workspaceTest[T, M <: TestUtil.BaseModule: Discover](m: M)
+  def workspaceTest[T, M <: TestUtil.BaseModule: Discover](m: M, resourcePath: Path = resourcePath)
                                                          (t: TestEvaluator[M] => T)
                                                          (implicit tp: TestPath): T = {
     val eval = new TestEvaluator(m)
@@ -300,11 +300,27 @@ object HelloWorldTests extends TestSuite {
           read(runResult) == "hello rockjam, your age is: 25"
         )
       }
-      'notRunWithoutMainClass - workspaceTest(HelloWorldWithoutMain){eval =>
-        val Left(Result.Exception(err, _)) = eval.apply(HelloWorldWithoutMain.core.run())
+      'notRunWithoutMainClass - workspaceTest(
+        HelloWorldWithoutMain,
+        pwd / 'scalalib / 'test / 'resources / "hello-world-no-main"
+      ){eval =>
+        val Left(Result.Failure(_, None)) = eval.apply(HelloWorldWithoutMain.core.run())
+      }
+
+      'runDiscoverMainClass - workspaceTest(HelloWorldWithoutMain){eval =>
+        // Make sure even if there isn't a main class defined explicitly, it gets
+        // discovered by Zinc and used
+        val runResult = eval.outPath / 'core / 'run / 'dest / "hello-mill"
+        val Right((_, evalCount)) = eval.apply(
+          HelloWorldWithoutMain.core.run(runResult.toString)
+        )
+
+        assert(evalCount > 0)
+
 
         assert(
-          err.isInstanceOf[RuntimeException]
+          exists(runResult),
+          read(runResult) == "hello rockjam, your age is: 25"
         )
       }
     }
@@ -337,12 +353,12 @@ object HelloWorldTests extends TestSuite {
           read(runResult) == "hello rockjam, your age is: 25"
         )
       }
-      'notRunWithoutMainClass - workspaceTest(HelloWorldWithoutMain){eval =>
-        val Left(Result.Exception(err, _)) = eval.apply(HelloWorldWithoutMain.core.runLocal())
+      'notRunWithoutMainClass - workspaceTest(
+        HelloWorldWithoutMain,
+        pwd / 'scalalib / 'test / 'resources / "hello-world-no-main"
+      ){eval =>
+        val Left(Result.Failure(_, None)) = eval.apply(HelloWorldWithoutMain.core.runLocal())
 
-        assert(
-          err.isInstanceOf[RuntimeException]
-        )
       }
     }
     'jar - {