From f682983bf95ce8e507d8b2c2e90614118359456f Mon Sep 17 00:00:00 2001
From: Cedric Beust <cedric@beust.com>
Date: Wed, 4 May 2016 23:26:01 -0800
Subject: [PATCH] Add ITestJvmFlagInterceptor.

---
 .../beust/kobalt/api/ITestJvmFlagInterceptor.kt | 14 ++++++++++++++
 .../com/beust/kobalt/internal/GenericRunner.kt  | 17 ++++++++++++-----
 .../beust/kobalt/internal/KobaltPluginXml.kt    |  6 ++++--
 3 files changed, 30 insertions(+), 7 deletions(-)
 create mode 100644 modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/ITestJvmFlagInterceptor.kt

diff --git a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/ITestJvmFlagInterceptor.kt b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/ITestJvmFlagInterceptor.kt
new file mode 100644
index 00000000..389f4704
--- /dev/null
+++ b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/api/ITestJvmFlagInterceptor.kt
@@ -0,0 +1,14 @@
+package com.beust.kobalt.api
+
+/**
+ * Plug-ins that add flags to the JVM used to run tests should implement this interface.
+ */
+interface ITestJvmFlagInterceptor : IInterceptor {
+    /**
+     * @return the list of all flags that should be used. If you only want to add flags to the current list,
+     * just return the concatenation of @param[currentFlags] and your own list (or use ITestJvmFlagContributor).
+     * If you actually alter the list of flags, make sure you don't remove anything critical from @param[currentFlags].
+     */
+    fun testJvmFlagsFor(project: Project, context: KobaltContext, currentFlags: List<String>) : List<String>
+}
+
diff --git a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/GenericRunner.kt b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/GenericRunner.kt
index 73011411..0fad304c 100644
--- a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/GenericRunner.kt
+++ b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/GenericRunner.kt
@@ -66,21 +66,28 @@ abstract class GenericTestRunner: ITestRunnerContributor {
                     add(mainClass)
                 }
 
-                // JVM args from the contributors
                 val pluginInfo = Kobalt.INJECTOR.getInstance(PluginInfo::class.java)
 
+                // JVM flags from the contributors
                 val flagsFromContributors = pluginInfo.testJvmFlagContributors.flatMap {
                     it.testJvmFlagsFor(project, context, jvmFlags)
                 }
 
-                if (flagsFromContributors.any()) {
-                    log(2, "Adding JVM flags from contributors: " + flagsFromContributors)
+                // JVM flags from the interceptors (these overwrite flags instead of just adding to the list)
+                var interceptedArgs = ArrayList(flagsFromContributors)
+                pluginInfo.testJvmFlagInterceptors.forEach {
+                    val newFlags = it.testJvmFlagsFor(project, context, interceptedArgs)
+                    interceptedArgs.clear()
+                    interceptedArgs.addAll(newFlags)
+                }
+
+                if (interceptedArgs.any()) {
+                    log(2, "Final JVM test flags after running the contributors and interceptors: $interceptedArgs")
                 }
 
                 val allArgs = arrayListOf<String>().apply {
                     add(java!!.absolutePath)
-                    addAll(flagsFromContributors)
-                    addAll(jvmFlags)
+                    addAll(interceptedArgs)
                     addAll(args)
                 }
 
diff --git a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/KobaltPluginXml.kt b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/KobaltPluginXml.kt
index e39a0e9f..f3138a75 100644
--- a/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/KobaltPluginXml.kt
+++ b/modules/kobalt-plugin-api/src/main/kotlin/com/beust/kobalt/internal/KobaltPluginXml.kt
@@ -80,6 +80,7 @@ class PluginInfo(val xml: KobaltPluginXml, val classLoader: ClassLoader?) {
     val buildConfigContributors = arrayListOf<IBuildConfigContributor>()
     val mavenIdInterceptors = arrayListOf<IMavenIdInterceptor>()
     val testJvmFlagContributors = arrayListOf<ITestJvmFlagContributor>()
+    val testJvmFlagInterceptors = arrayListOf<ITestJvmFlagInterceptor>()
 
     companion object {
         /**
@@ -166,6 +167,7 @@ class PluginInfo(val xml: KobaltPluginXml, val classLoader: ClassLoader?) {
 
                 // Not documented yet
                 if (this is ITestJvmFlagContributor) testJvmFlagContributors.add(this)
+                if (this is ITestJvmFlagInterceptor) testJvmFlagInterceptors.add(this)
             }
         }
     }
@@ -178,7 +180,7 @@ class PluginInfo(val xml: KobaltPluginXml, val classLoader: ClassLoader?) {
                 compilerContributors, docContributors, sourceDirContributors,
                 testSourceDirContributors, buildConfigFieldContributors,
                 taskContributors, assemblyContributors,
-                incrementalAssemblyContributors, testJvmFlagContributors
+                incrementalAssemblyContributors, testJvmFlagInterceptors
             ).forEach {
                 it.forEach {
                     it.cleanUpActors()
@@ -214,7 +216,7 @@ class PluginInfo(val xml: KobaltPluginXml, val classLoader: ClassLoader?) {
         buildConfigContributors.addAll(pluginInfo.buildConfigContributors)
         assemblyContributors.addAll(pluginInfo.assemblyContributors)
         incrementalAssemblyContributors.addAll(pluginInfo.incrementalAssemblyContributors)
-        testJvmFlagContributors.addAll(pluginInfo.testJvmFlagContributors)
+        testJvmFlagInterceptors.addAll(pluginInfo.testJvmFlagInterceptors)
     }
 }