fix #233 add append and exclude rules to assembly (#309)

* fix #233 add append and exclude rules to assembly

* handle existing files and concatenation when file already exists in assembly

* add assembly tests for append rules

* tests for append patterns

* tests for exclude patterns

* make append algorithm use single map with fold over classpathIterator

* move assembly rules logic to method

* move grouping method to Assembly object, make assemblyRules Seq[_] rather than T[Seq[_]]

* add test cases for when there are no rules

* keep default parameter in createAssembly not to break CI

* add one more reference.conf entry to tests

7 files changed, 295 insertions, 16 deletions

diff --git a/scalalib/src/mill/scalalib/JavaModule.scala b/scalalib/src/mill/scalalib/JavaModule.scala
index 26c58a8b..79fad9aa 100644
--- a/scalalib/src/mill/scalalib/JavaModule.scala
+++ b/scalalib/src/mill/scalalib/JavaModule.scala
@@ -7,7 +7,7 @@ import coursier.Repository
 import mill.define.Task
 import mill.define.TaskModule
 import mill.eval.{PathRef, Result}
-import mill.modules.Jvm
+import mill.modules.{Assembly, Jvm}
 import mill.modules.Jvm.{createAssembly, createJar}
 import Lib._
 import mill.scalalib.publish.{Artifact, Scope}
@@ -108,6 +108,8 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     }
   }
 
+  def assemblyRules: Seq[Assembly.Rule] = Assembly.defaultRules
+
   def sources = T.sources{ millSourcePath / 'src }
   def resources = T.sources{ millSourcePath / 'resources }
   def generatedSources = T{ Seq.empty[PathRef] }
@@ -130,6 +132,7 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
       upstreamCompileOutput()
     )
   }
+
   def localClasspath = T{
     resources() ++ Agg(compile().classes)
   }
@@ -158,7 +161,11 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     * upstream dependencies do not change
     */
   def upstreamAssembly = T{
-    createAssembly(upstreamAssemblyClasspath().map(_.path), mainClass())
+    createAssembly(
+      upstreamAssemblyClasspath().map(_.path),
+      mainClass(),
+      assemblyRules = assemblyRules
+    )
   }
 
   def assembly = T{
@@ -166,11 +173,11 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
       Agg.from(localClasspath().map(_.path)),
       mainClass(),
       prependShellScript(),
-      Some(upstreamAssembly().path)
+      Some(upstreamAssembly().path),
+      assemblyRules
     )
   }
 
-
   def jar = T{
     createJar(
       localClasspath().map(_.path).filter(exists),
diff --git a/scalalib/test/resources/hello-world-multi/core/resources/reference.conf b/scalalib/test/resources/hello-world-multi/core/resources/reference.conf
new file mode 100644
index 00000000..afb44467
--- /dev/null
+++ b/scalalib/test/resources/hello-world-multi/core/resources/reference.conf
@@ -0,0 +1,4 @@
+##############################
+# Core Reference Config File #
+##############################
+bar.baz=hello
diff --git a/scalalib/test/resources/hello-world-multi/core/src/Main.scala b/scalalib/test/resources/hello-world-multi/core/src/Main.scala
new file mode 100644
index 00000000..5cbb75cf
--- /dev/null
+++ b/scalalib/test/resources/hello-world-multi/core/src/Main.scala
@@ -0,0 +1,5 @@
+object Main extends App {
+  val person = Person.fromString("rockjam:25")
+  println(s"hello ${person.name}, your age is: ${person.age}")
+}
+
diff --git a/scalalib/test/resources/hello-world-multi/model/resources/reference.conf b/scalalib/test/resources/hello-world-multi/model/resources/reference.conf
new file mode 100644
index 00000000..8e7ed298
--- /dev/null
+++ b/scalalib/test/resources/hello-world-multi/model/resources/reference.conf
@@ -0,0 +1,4 @@
+###############################
+# Model Reference Config File #
+###############################
+foo.bar=2
diff --git a/scalalib/test/resources/hello-world-multi/model/src/Person.scala b/scalalib/test/resources/hello-world-multi/model/src/Person.scala
new file mode 100644
index 00000000..23e3821e
--- /dev/null
+++ b/scalalib/test/resources/hello-world-multi/model/src/Person.scala
@@ -0,0 +1,8 @@
+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/resources/hello-world/core/resources/reference.conf b/scalalib/test/resources/hello-world/core/resources/reference.conf
new file mode 100644
index 00000000..bf0f66ae
--- /dev/null
+++ b/scalalib/test/resources/hello-world/core/resources/reference.conf
@@ -0,0 +1,4 @@
+########################################
+# My application Reference Config File #
+########################################
+akka.http.client.user-agent-header="hello-world-client"
diff --git a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
index 46cf82de..81a5b8a6 100644
--- a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
+++ b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
@@ -3,15 +3,14 @@ package mill.scalalib
 import java.util.jar.JarFile
 
 import ammonite.ops._
-import ammonite.ops.ImplicitWd._
 import mill._
-import mill.define.{Discover, Target}
+import mill.define.Target
 import mill.eval.{Evaluator, Result}
+import mill.modules.Assembly
 import mill.scalalib.publish._
 import mill.util.{TestEvaluator, TestUtil}
 import mill.scalalib.publish.VersionControl
 import utest._
-
 import utest.framework.TestPath
 
 import scala.collection.JavaConverters._
@@ -21,10 +20,14 @@ object HelloWorldTests extends TestSuite {
   trait HelloBase extends TestUtil.BaseModule{
     def millSourcePath =  TestUtil.getSrcPathBase() / millOuterCtx.enclosing.split('.')
   }
+
   trait HelloWorldModule extends scalalib.ScalaModule {
     def scalaVersion = "2.12.4"
   }
 
+  trait HelloWorldModuleWithMain extends HelloWorldModule {
+    def mainClass = Some("Main")
+  }
 
   object HelloWorld extends HelloBase {
     object core extends HelloWorldModule
@@ -45,11 +48,96 @@ object HelloWorldTests extends TestSuite {
   }
 
   object HelloWorldWithMain extends HelloBase {
-    object core extends HelloWorldModule{
-      def mainClass = Some("Main")
+    object core extends HelloWorldModuleWithMain
+  }
+
+  val akkaHttpDeps = Agg(ivy"com.typesafe.akka::akka-http:10.0.13")
+
+  object HelloWorldAkkaHttpAppend extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def ivyDeps = akkaHttpDeps
+
+      def assemblyRules = Seq(Assembly.Rule.Append("reference.conf"))
     }
   }
 
+  object HelloWorldAkkaHttpExclude extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def ivyDeps = akkaHttpDeps
+
+      def assemblyRules = Seq(Assembly.Rule.Exclude("reference.conf"))
+    }
+  }
+
+  object HelloWorldAkkaHttpAppendPattern extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def ivyDeps = akkaHttpDeps
+
+      def assemblyRules = Seq(Assembly.Rule.AppendPattern(".*.conf"))
+    }
+  }
+
+  object HelloWorldAkkaHttpExcludePattern extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def ivyDeps = akkaHttpDeps
+
+      def assemblyRules = Seq(Assembly.Rule.ExcludePattern(".*.conf"))
+    }
+  }
+
+  object HelloWorldAkkaHttpNoRules extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def ivyDeps = akkaHttpDeps
+
+      def assemblyRules = Seq.empty
+    }
+  }
+
+  object HelloWorldMultiAppend extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def moduleDeps = Seq(model)
+
+      def assemblyRules = Seq(Assembly.Rule.Append("reference.conf"))
+    }
+    object model extends HelloWorldModule
+  }
+
+  object HelloWorldMultiExclude extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def moduleDeps = Seq(model)
+
+      def assemblyRules = Seq(Assembly.Rule.Exclude("reference.conf"))
+    }
+    object model extends HelloWorldModule
+  }
+
+  object HelloWorldMultiAppendPattern extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def moduleDeps = Seq(model)
+
+      def assemblyRules = Seq(Assembly.Rule.AppendPattern(".*.conf"))
+    }
+    object model extends HelloWorldModule
+  }
+
+  object HelloWorldMultiExcludePattern extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def moduleDeps = Seq(model)
+
+      def assemblyRules = Seq(Assembly.Rule.ExcludePattern(".*.conf"))
+    }
+    object model extends HelloWorldModule
+  }
+
+  object HelloWorldMultiNoRules extends HelloBase {
+    object core extends HelloWorldModuleWithMain {
+      def moduleDeps = Seq(model)
+
+      def assemblyRules = Seq.empty
+    }
+    object model extends HelloWorldModule
+  }
+
   object HelloWorldWarnUnused extends HelloBase{
     object core extends HelloWorldModule {
       def scalacOptions = T(Seq("-Ywarn-unused"))
@@ -158,6 +246,15 @@ object HelloWorldTests extends TestSuite {
     attrs.get(Name.MAIN_CLASS).map(_.asInstanceOf[String])
   }
 
+  def jarEntries(jar: JarFile): Set[String] = {
+    jar.entries().asScala.map(_.getName).toSet
+  }
+
+  def readFileFromJar(jar: JarFile, name: String): String = {
+    val is = jar.getInputStream(jar.getEntry(name))
+    read(is)
+  }
+
   def compileClassfiles = Seq[RelPath](
     "Main.class",
     "Main$.class",
@@ -178,14 +275,14 @@ object HelloWorldTests extends TestSuite {
     cp(resourcePath, m.millSourcePath)
     t(eval)
   }
-  
+
 
 
 
   def tests: Tests = Tests {
     'scalaVersion - {
-      
-      'fromBuild - workspaceTest(HelloWorld){eval => 
+
+      'fromBuild - workspaceTest(HelloWorld){eval =>
         val Right((result, evalCount)) = eval.apply(HelloWorld.core.scalaVersion)
 
         assert(
@@ -421,10 +518,11 @@ object HelloWorldTests extends TestSuite {
         val jarFile = new JarFile(result.path.toIO)
         val entries = jarFile.entries().asScala.map(_.getName).toSet
 
-        val manifestFiles = Seq[RelPath](
-          "META-INF" / "MANIFEST.MF"
+        val otherFiles = Seq[RelPath](
+          "META-INF" / "MANIFEST.MF",
+          "reference.conf"
         )
-        val expectedFiles = compileClassfiles ++ manifestFiles
+        val expectedFiles = compileClassfiles ++ otherFiles
 
         assert(
           entries.nonEmpty,
@@ -452,7 +550,7 @@ object HelloWorldTests extends TestSuite {
           evalCount > 0
         )
         val jarFile = new JarFile(result.path.toIO)
-        val entries = jarFile.entries().asScala.map(_.getName).toSet
+        val entries = jarEntries(jarFile)
 
         val mainPresent = entries.contains("Main.class")
         assert(mainPresent)
@@ -462,6 +560,155 @@ object HelloWorldTests extends TestSuite {
         assert(mainClass.contains("Main"))
       }
 
+      'assemblyRules - {
+        def checkAppend[M <: TestUtil.BaseModule](module: M,
+                                                  target: Target[PathRef]) =
+          workspaceTest(module) { eval =>
+            val Right((result, _)) = eval.apply(target)
+
+            val jarFile = new JarFile(result.path.toIO)
+
+            assert(jarEntries(jarFile).contains("reference.conf"))
+
+            val referenceContent = readFileFromJar(jarFile, "reference.conf")
+
+            assert(
+              // akka modules configs are present
+              referenceContent.contains("akka-http Reference Config File"),
+              referenceContent.contains("akka-http-core Reference Config File"),
+              referenceContent.contains("Akka Actor Reference Config File"),
+              referenceContent.contains("Akka Stream Reference Config File"),
+              // our application config is present too
+              referenceContent.contains("My application Reference Config File"),
+              referenceContent.contains(
+                """akka.http.client.user-agent-header="hello-world-client""""
+              )
+            )
+          }
+
+        val helloWorldMultiResourcePath = pwd / 'scalalib / 'test / 'resources / "hello-world-multi"
+
+        def checkAppendMulti[M <: TestUtil.BaseModule](
+            module: M,
+            target: Target[PathRef]) =
+          workspaceTest(
+            module,
+            resourcePath = helloWorldMultiResourcePath
+          ) { eval =>
+            val Right((result, _)) = eval.apply(target)
+
+            val jarFile = new JarFile(result.path.toIO)
+
+            assert(jarEntries(jarFile).contains("reference.conf"))
+
+            val referenceContent = readFileFromJar(jarFile, "reference.conf")
+
+            assert(
+              // reference config from core module
+              referenceContent.contains("Core Reference Config File"),
+              // reference config from model module
+              referenceContent.contains("Model Reference Config File"),
+              // concatenated content
+              referenceContent.contains("bar.baz=hello"),
+              referenceContent.contains("foo.bar=2")
+            )
+          }
+
+        'appendWithDeps - checkAppend(
+          HelloWorldAkkaHttpAppend,
+          HelloWorldAkkaHttpAppend.core.assembly
+        )
+        'appendMultiModule - checkAppendMulti(
+          HelloWorldMultiAppend,
+          HelloWorldMultiAppend.core.assembly
+        )
+        'appendPatternWithDeps - checkAppend(
+          HelloWorldAkkaHttpAppendPattern,
+          HelloWorldAkkaHttpAppendPattern.core.assembly
+        )
+        'appendPatternMultiModule - checkAppendMulti(
+          HelloWorldMultiAppendPattern,
+          HelloWorldMultiAppendPattern.core.assembly
+        )
+
+        def checkExclude[M <: TestUtil.BaseModule](module: M,
+                                                   target: Target[PathRef],
+                                                   resourcePath: Path = resourcePath
+                                                  ) =
+          workspaceTest(module, resourcePath) { eval =>
+            val Right((result, _)) = eval.apply(target)
+
+            val jarFile = new JarFile(result.path.toIO)
+
+            assert(!jarEntries(jarFile).contains("reference.conf"))
+          }
+
+        'excludeWithDeps - checkExclude(
+          HelloWorldAkkaHttpExclude,
+          HelloWorldAkkaHttpExclude.core.assembly
+        )
+        'excludeMultiModule - checkExclude(
+          HelloWorldMultiExclude,
+          HelloWorldMultiExclude.core.assembly,
+          resourcePath = helloWorldMultiResourcePath
+
+        )
+        'excludePatternWithDeps - checkExclude(
+          HelloWorldAkkaHttpExcludePattern,
+          HelloWorldAkkaHttpExcludePattern.core.assembly
+        )
+        'excludePatternMultiModule - checkExclude(
+          HelloWorldMultiExcludePattern,
+          HelloWorldMultiExcludePattern.core.assembly,
+          resourcePath = helloWorldMultiResourcePath
+        )
+
+        'writeFirstWhenNoRule - {
+          'withDeps - workspaceTest(HelloWorldAkkaHttpNoRules) { eval =>
+            val Right((result, _)) = eval.apply(HelloWorldAkkaHttpNoRules.core.assembly)
+
+            val jarFile = new JarFile(result.path.toIO)
+
+            assert(jarEntries(jarFile).contains("reference.conf"))
+
+            val referenceContent = readFileFromJar(jarFile, "reference.conf")
+
+            val allOccurrences = Seq(
+              referenceContent.contains("akka-http Reference Config File"),
+              referenceContent.contains("akka-http-core Reference Config File"),
+              referenceContent.contains("Akka Actor Reference Config File"),
+              referenceContent.contains("Akka Stream Reference Config File"),
+              referenceContent.contains("My application Reference Config File")
+            )
+
+            val timesOcccurres = allOccurrences.find(identity).size
+
+            assert(timesOcccurres == 1)
+          }
+
+          'multiModule - workspaceTest(
+            HelloWorldMultiNoRules,
+            resourcePath = helloWorldMultiResourcePath
+          ) { eval =>
+            val Right((result, _)) = eval.apply(HelloWorldMultiNoRules.core.assembly)
+
+            val jarFile = new JarFile(result.path.toIO)
+
+            assert(jarEntries(jarFile).contains("reference.conf"))
+
+            val referenceContent = readFileFromJar(jarFile, "reference.conf")
+
+            assert(
+              referenceContent.contains("Model Reference Config File"),
+              referenceContent.contains("foo.bar=2"),
+
+              !referenceContent.contains("Core Reference Config File"),
+              !referenceContent.contains("bar.baz=hello")
+            )
+          }
+        }
+      }
+
       'run - workspaceTest(HelloWorldWithMain){eval =>
         val Right((result, evalCount)) = eval.apply(HelloWorldWithMain.core.assembly)