feat: Added option for default test time and average time for smart shard

README.md

@@ -169,6 +169,14 @@ flank:
   ## Default: false
   # smart-flank-disable-upload: false
 
+  ## Enable using average time from previous tests duration when using SmartShard and tests did not run before.
+  ## Default: false
+  # use-average-test-time-for-new-tests: true
+
+  ## Set default test time used for calculating shards.
+  ## Default: 120.0
+  # default-test-time: 15
+
   ## Disables sharding. Useful for parameterized tests.
   # disable-sharding: false
 
@@ -385,6 +393,14 @@ flank:
   ## Default: false
   # smart-flank-disable-upload: false
 
+  ## Enable using average time from previous tests duration when using SmartShard and tests did not run before.
+  ## Default: false
+  # use-average-test-time-for-new-tests: true
+
+  ## Set default test time used for calculating shards.
+  ## Default: 120.0
+  # default-test-time: 15
+
   ## Disables sharding. Useful for parameterized tests.
   # disable-sharding: false
 

test_runner/flank.ios.yml

@@ -101,6 +101,14 @@ flank:
   ## Default: false
   # smart-flank-disable-upload: false
 
+  ## Enable using average time from previous tests duration when using SmartShard and tests did not run before.
+  ## Default: false
+  # use-average-test-time-for-new-tests: true
+
+  ## Set default test time used for calculating shards.
+  ## Default: 120.0
+  # default-test-time: 15
+
   ## Disables sharding. Useful for parameterized tests.
   # disable-sharding: false
 

test_runner/flank.yml

@@ -161,6 +161,14 @@ flank:
   ## Default: false
   # smart-flank-disable-upload: false
 
+  ## Enable using average time from previous tests duration when using SmartShard and tests did not run before.
+  ## Default: false
+  # use-average-test-time-for-new-tests: true
+
+  ## Set default test time used for calculating shards.
+  ## Default: 120.0
+  # default-test-time: 15
+
   ## Disables sharding. Useful for parameterized tests.
   # disable-sharding: false
 

test_runner/src/main/kotlin/ftl/args/AndroidArgs.kt

@@ -58,6 +58,8 @@ AndroidArgs
       num-test-runs: $repeatTests
       smart-flank-gcs-path: $smartFlankGcsPath
       smart-flank-disable-upload: $smartFlankDisableUpload
+      default-test-time: $defaultTestTime
+      use-average-test-time-for-new-tests: $useAverageTestTimeForNewTests
       files-to-download:${ArgsToString.listToString(filesToDownload)}
       test-targets-always-run:${ArgsToString.listToString(testTargetsAlwaysRun)}
       disable-sharding: $disableSharding

test_runner/src/main/kotlin/ftl/args/CommonArgs.kt

@@ -34,5 +34,7 @@ data class CommonArgs(
     override val ignoreFailedTests: Boolean,
     override val keepFilePath: Boolean,
     override val outputStyle: OutputStyle,
-    override val disableResultsUpload: Boolean
+    override val disableResultsUpload: Boolean,
+    override val defaultTestTime: Double,
+    override val useAverageTestTimeForNewTests: Boolean,
 ) : IArgs

test_runner/src/main/kotlin/ftl/args/CreateCommonArgs.kt

@@ -41,7 +41,9 @@ fun CommonConfig.createCommonArgs(
     filesToDownload = flank.filesToDownload!!,
     disableSharding = flank.disableSharding!!,
     localResultDir = flank.localResultsDir!!,
-    disableResultsUpload = flank.disableResultsUpload!!
+    disableResultsUpload = flank.disableResultsUpload!!,
+    defaultTestTime = flank.defaultTestTime!!,
+    useAverageTestTimeForNewTests = flank.useAverageTestTimeForNewTests!!
 ).apply {
     ArgsHelper.createJunitBucket(project, smartFlankGcsPath)
 }

test_runner/src/main/kotlin/ftl/args/IArgs.kt

@@ -59,6 +59,9 @@ interface IArgs {
     val inVirtualRange: Boolean
         get() = maxTestShards in AVAILABLE_VIRTUAL_SHARD_COUNT_RANGE
 
+    val defaultTestTime: Double
+    val useAverageTestTimeForNewTests: Boolean
+
     fun useLocalResultDir() = localResultDir != defaultLocalResultsDir
 
     companion object {

test_runner/src/main/kotlin/ftl/args/IosArgs.kt

@@ -43,6 +43,8 @@ IosArgs
       num-test-runs: $repeatTests
       smart-flank-gcs-path: $smartFlankGcsPath
       smart-flank-disable-upload: $smartFlankDisableUpload
+      default-test-time: $defaultTestTime
+      use-average-test-time-for-new-tests: $useAverageTestTimeForNewTests
       test-targets-always-run:${ArgsToString.listToString(testTargetsAlwaysRun)}
       files-to-download:${ArgsToString.listToString(filesToDownload)}
       keep-file-path: $keepFilePath

test_runner/src/main/kotlin/ftl/config/common/CommonFlankConfig.kt

@@ -7,6 +7,7 @@ import ftl.args.ArgsHelper
 import ftl.args.yml.IYmlKeys
 import ftl.config.Config
 import ftl.config.FtlConstants
+import ftl.shard.DEFAULT_TEST_TIME_SEC
 import picocli.CommandLine
 
 /** Flank specific parameters for both iOS and Android */
@@ -140,6 +141,20 @@ data class CommonFlankConfig @JsonIgnore constructor(
     @set:JsonProperty("disable-results-upload")
     var disableResultsUpload: Boolean? by data
 
+    @set:CommandLine.Option(
+        names = ["--default-test-time"],
+        description = ["Set default test time used for calculating shards."]
+    )
+    @set:JsonProperty("default-test-time")
+    var defaultTestTime: Double? by data
+
+    @set:CommandLine.Option(
+        names = ["--use-average-test-time-for-new-tests"],
+        description = ["Enable using average time from previous tests duration when using SmartShard and tests did not run before."]
+    )
+    @set:JsonProperty("use-average-test-time-for-new-tests")
+    var useAverageTestTimeForNewTests: Boolean? by data
+
     constructor() : this(mutableMapOf<String, Any?>().withDefault { null })
 
     companion object : IYmlKeys {
@@ -184,6 +199,8 @@ data class CommonFlankConfig @JsonIgnore constructor(
             keepFilePath = false
             outputStyle = null
             disableResultsUpload = false
+            defaultTestTime = DEFAULT_TEST_TIME_SEC
+            useAverageTestTimeForNewTests = false
         }
     }
 }

test_runner/src/main/kotlin/ftl/shard/Shard.kt

@@ -45,7 +45,7 @@ fun createShardsByShardCount(
     val maxShards = maxShards(args.maxTestShards, forcedShardCount)
 
     val previousMethodDurations = createTestMethodDurationMap(oldTestResult, args)
-    val testCases = createTestCases(testsToRun, previousMethodDurations)
+    val testCases = createTestCases(testsToRun, previousMethodDurations, args)
         .sortedByDescending(TestMethod::time) // We want to iterate over testcase going from slowest to fastest
 
     val testCount = getNumberOfNotIgnoredTestCases(testCases)

test_runner/src/main/kotlin/ftl/shard/ShardCount.kt

@@ -26,14 +26,26 @@ private fun calculateShardCount(
     testsToRun: List<FlankTestMethod>,
     oldTestResult: JUnitTestResult,
     args: IArgs
-): Int = calculateShardCount(
-    args = args,
-    testsTotalTime = testTotalTime(testsToRun, createTestMethodDurationMap(oldTestResult, args)),
-    testsToRunCount = testsToRun.size
-)
-
-private fun testTotalTime(testsToRun: List<FlankTestMethod>, previousMethodDurations: Map<String, Double>): Double =
-    testsToRun.sumByDouble { flankTestMethod -> getTestMethodTime(flankTestMethod, previousMethodDurations) }
+): Int {
+    val previousMethodDurations = createTestMethodDurationMap(oldTestResult, args)
+    return calculateShardCount(
+        args = args,
+        testsTotalTime = testTotalTime(testsToRun, previousMethodDurations, args.fallbackTestTime(previousMethodDurations)),
+        testsToRunCount = testsToRun.size
+    )
+}
+
+private fun testTotalTime(
+    testsToRun: List<FlankTestMethod>,
+    previousMethodDurations: Map<String, Double>,
+    defaultTestTime: Double
+) = testsToRun.sumByDouble { flankTestMethod ->
+    getTestMethodTime(
+        flankTestMethod,
+        previousMethodDurations,
+        defaultTestTime
+    )
+}
 
 private fun calculateShardCount(
     args: IArgs,

test_runner/src/main/kotlin/ftl/shard/TestCasesCreator.kt

@@ -1,12 +1,28 @@
 package ftl.shard
 
+import ftl.args.IArgs
 import ftl.util.FlankTestMethod
 
 data class TestMethod(
     val name: String,
     val time: Double
 )
 
-fun createTestCases(testsToRun: List<FlankTestMethod>, previousMethodDurations: Map<String, Double>): List<TestMethod> {
-    return testsToRun.map { TestMethod(it.testName, getTestMethodTime(it, previousMethodDurations)) }
+fun createTestCases(
+    testsToRun: List<FlankTestMethod>,
+    previousMethodDurations: Map<String, Double>,
+    args: IArgs,
+): List<TestMethod> {
+    val defaultTestTime = args.fallbackTestTime(previousMethodDurations)
+    return testsToRun
+        .map {
+        TestMethod(
+            name = it.testName,
+            time = getTestMethodTime(
+                flankTestMethod = it,
+                previousMethodDurations = previousMethodDurations,
+                defaultTestTime = defaultTestTime
+            )
+        )
+    }
 }

test_runner/src/main/kotlin/ftl/shard/TestMethodTime.kt

@@ -1,18 +1,26 @@
 package ftl.shard
 
-import com.google.common.annotations.VisibleForTesting
+import ftl.args.IArgs
 import ftl.util.FlankTestMethod
 
-// When a test does not have previous results to reference, fall back to this run time.
-@VisibleForTesting
-const val DEFAULT_TEST_TIME_SEC = 120.0
+fun getTestMethodTime(
+    flankTestMethod: FlankTestMethod,
+    previousMethodDurations: Map<String, Double>,
+    defaultTestTime: Double
+) = if (flankTestMethod.ignored) IGNORE_TEST_TIME else previousMethodDurations.getOrElse(flankTestMethod.testName) {
+    defaultTestTime
+}
+
+fun IArgs.fallbackTestTime(
+    previousMethodDurations: Map<String, Double>
+) = if (useAverageTestTimeForNewTests) previousMethodDurations.averageTestTime(defaultTestTime) else defaultTestTime
+
+private fun Map<String, Double>.averageTestTime(defaultTestTime: Double) = values
+    .filter { it > IGNORE_TEST_TIME }
+    .average()
+    .takeIf { !it.isNaN() }
+    ?: defaultTestTime
 
-@VisibleForTesting
 const val IGNORE_TEST_TIME = 0.0
 
-fun getTestMethodTime(flankTestMethod: FlankTestMethod, previousMethodDurations: Map<String, Double>): Double {
-    return if (flankTestMethod.ignored) IGNORE_TEST_TIME else previousMethodDurations.getOrDefault(
-        flankTestMethod.testName,
-        DEFAULT_TEST_TIME_SEC
-    )
-}
+const val DEFAULT_TEST_TIME_SEC = 120.0

test_runner/src/test/kotlin/ftl/args/AndroidArgsTest.kt

@@ -108,6 +108,8 @@ class AndroidArgsTest {
           max-test-shards: 7
           shard-time: 60
           num-test-runs: 8
+          default-test-time: 15.0
+          use-average-test-time-for-new-tests: true
           files-to-download:
             - /sdcard/screenshots
             - /sdcard/screenshots2
@@ -317,6 +319,8 @@ AndroidArgs
       num-test-runs: 8
       smart-flank-gcs-path:${' '}
       smart-flank-disable-upload: false
+      default-test-time: 15.0
+      use-average-test-time-for-new-tests: true
       files-to-download:
         - /sdcard/screenshots
         - /sdcard/screenshots2
@@ -382,6 +386,8 @@ AndroidArgs
       num-test-runs: 1
       smart-flank-gcs-path: 
       smart-flank-disable-upload: false
+      default-test-time: 120.0
+      use-average-test-time-for-new-tests: false
       files-to-download:
       test-targets-always-run:
       disable-sharding: false
@@ -1676,6 +1682,61 @@ AndroidArgs
         """.trimIndent()
         AndroidArgs.load(yaml)
     }
+
+    @Test
+    fun `should set defaultTestTime`() {
+        val yaml = """
+        gcloud:
+          app: $appApk
+          test: $testApk
+        flank:
+          max-test-shards: -1
+          default-test-time: 15
+        """.trimIndent()
+        val args = AndroidArgs.load(yaml)
+        assertEquals(args.defaultTestTime, 15.0, 0.01)
+    }
+
+    @Test
+    fun `should set defaultTestTime to default value if not specified`() {
+        val yaml = """
+        gcloud:
+          app: $appApk
+          test: $testApk
+        flank:
+          max-test-shards: -1
+          use-average-test-time-for-new-tests: true
+        """.trimIndent()
+        val args = AndroidArgs.load(yaml)
+        assertEquals(args.defaultTestTime, 120.0, 0.01)
+    }
+
+    @Test
+    fun `should useAverageTestTimeForNewTests set to true`() {
+        val yaml = """
+        gcloud:
+          app: $appApk
+          test: $testApk
+        flank:
+          max-test-shards: -1
+          use-average-test-time-for-new-tests: true
+        """.trimIndent()
+        val args = AndroidArgs.load(yaml)
+        assertTrue(args.useAverageTestTimeForNewTests)
+    }
+
+    @Test
+    fun `should useAverageTestTimeForNewTests set to false by defaul`() {
+        val yaml = """
+        gcloud:
+          app: $appApk
+          test: $testApk
+        flank:
+          max-test-shards: -1
+        """.trimIndent()
+        val args = AndroidArgs.load(yaml)
+        assertFalse(args.useAverageTestTimeForNewTests)
+    }
 }
 
 private fun AndroidArgs.Companion.load(yamlData: String, cli: AndroidRunCommand? = null): AndroidArgs =