Add unit test for main method discovery

2 files changed, 19 insertions, 13 deletions

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]
 }