Issue #399: Smart Flank - Dynamic bucket counts v2 (#471)

* Issue #399: Smart Flank - Dynamic bucket counts

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

@@ -65,6 +65,7 @@ class AndroidArgs(
 
     private val flank = flankYml.flank
     override val testShards = cli?.testShards ?: flank.testShards
+    override val shardTime = cli?.shardTime ?: flank.shardTime
     override val repeatTests = cli?.repeatTests ?: flank.repeatTests
     override val smartFlankGcsPath = flank.smartFlankGcsPath
     override val testTargetsAlwaysRun = cli?.testTargetsAlwaysRun ?: flank.testTargetsAlwaysRun
@@ -160,6 +161,7 @@ ${devicesToString(devices)}
 
     flank:
       testShards: $testShards
+      shardTime: $shardTime
       repeatTests: $repeatTests
       smartFlankGcsPath: $smartFlankGcsPath
       files-to-download:

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

@@ -178,9 +178,12 @@ object ArgsHelper {
 
     fun calculateShards(filteredTests: List<String>, args: IArgs): List<List<String>> {
         val oldTestResult = GcStorage.downloadJunitXml(args) ?: JUnitTestResult(mutableListOf())
-        val shardsByTime = Shard.calculateShardsByTime(filteredTests, oldTestResult, args)
 
-        return testMethodsAlwaysRun(shardsByTime.stringShards(), args)
+        val shardCount = Shard.shardCountByTime(filteredTests, oldTestResult, args)
+
+        val shards = Shard.createShardsByShardCount(filteredTests, oldTestResult, args, shardCount)
+
+        return testMethodsAlwaysRun(shards.stringShards(), args)
     }
 
     private fun testMethodsAlwaysRun(shards: StringShards, args: IArgs): StringShards {

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

@@ -16,6 +16,7 @@ interface IArgs {
 
     // FlankYml
     val testShards: Int
+    val shardTime: Int
     val repeatTests: Int
     val smartFlankGcsPath: String
     val testTargetsAlwaysRun: List<String>

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

@@ -50,6 +50,7 @@ class IosArgs(
 
     private val flank = flankYml.flank
     override val testShards = cli?.testShards ?: flank.testShards
+    override val shardTime = cli?.shardTime ?: flank.shardTime
     override val repeatTests = cli?.repeatTests ?: flank.repeatTests
     override val smartFlankGcsPath = flank.smartFlankGcsPath
     override val testTargetsAlwaysRun = cli?.testTargetsAlwaysRun ?: flank.testTargetsAlwaysRun
@@ -122,6 +123,7 @@ ${devicesToString(devices)}
 
     flank:
       testShards: $testShards
+      shardTime: $shardTime
       repeatTests: $repeatTests
       smartFlankGcsPath: $smartFlankGcsPath
       test-targets-always-run:

test_runner/src/main/kotlin/ftl/args/yml/FlankYml.kt

@@ -9,6 +9,7 @@ import ftl.util.Utils.fatalError
 @JsonIgnoreProperties(ignoreUnknown = true)
 class FlankYmlParams(
     val testShards: Int = 1,
+    val shardTime: Int = -1,
     val repeatTests: Int = 1,
     val smartFlankGcsPath: String = "",
 
@@ -19,11 +20,14 @@ class FlankYmlParams(
     val filesToDownload: List<String> = emptyList()
 ) {
     companion object : IYmlKeys {
-        override val keys = listOf("testShards", "repeatTests", "smartFlankGcsPath", "test-targets-always-run", "files-to-download")
+        override val keys = listOf(
+            "testShards", "shardTime", "repeatTests", "smartFlankGcsPath", "test-targets-always-run", "files-to-download"
+        )
     }
 
     init {
         if (testShards <= 0 && testShards != -1) fatalError("testShards must be >= 1 or -1")
+        if (shardTime <= 0 && shardTime != -1) fatalError("shardTime must be >= 1 or -1")
         if (repeatTests < 1) fatalError("repeatTests must be >= 1")
 
         if (smartFlankGcsPath.isNotEmpty()) {

test_runner/src/main/kotlin/ftl/cli/firebase/test/android/AndroidRunCommand.kt

@@ -221,6 +221,12 @@ class AndroidRunCommand : Runnable {
     )
     var testShards: Int? = null
 
+    @Option(
+        names = ["--shard-time"],
+        description = ["The max amount of seconds each shard should run."]
+    )
+    var shardTime: Int? = null
+
     @Option(
         names = ["--repeat-tests"],
         description = ["The amount of times to repeat the test executions."]

test_runner/src/main/kotlin/ftl/cli/firebase/test/ios/IosRunCommand.kt

@@ -105,6 +105,12 @@ class IosRunCommand : Runnable {
     )
     var testShards: Int? = null
 
+    @Option(
+        names = ["--shard-time"],
+        description = ["The max amount of seconds each shard should run."]
+    )
+    var shardTime: Int? = null
+
     @Option(
         names = ["--repeat-tests"],
         description = ["The amount of times to repeat the test executions."]

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

@@ -52,13 +52,31 @@ object Shard {
         return "$classname/$testName"
     }
 
-    // take in the XML with timing info then return list of shards
-    fun calculateShardsByTime(
+    // take in the XML with timing info then return the shard count based on execution time
+    fun shardCountByTime(
         testsToRun: List<String>,
         oldTestResult: JUnitTestResult,
         args: IArgs
+    ): Int {
+        if (args.shardTime == -1) return -1
+
+        val junitMap = createJunitMap(oldTestResult, args)
+        val testsTotalTime = testsToRun.sumByDouble { junitMap[it] ?: 10.0 }
+
+        val shardsByTime = Math.ceil(testsTotalTime / args.shardTime).toInt()
+
+        // We need to respect the testShards
+        return Math.min(shardsByTime, args.testShards)
+    }
+
+    // take in the XML with timing info then return list of shards based on the amount of shards to use
+    fun createShardsByShardCount(
+        testsToRun: List<String>,
+        oldTestResult: JUnitTestResult,
+        args: IArgs,
+        forcedShardCount: Int = -1
     ): List<TestShard> {
-        val maxShards = args.testShards
+        val maxShards = if (forcedShardCount == -1) args.testShards else forcedShardCount
         val junitMap = createJunitMap(oldTestResult, args)
 
         var cacheMiss = 0

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

@@ -59,6 +59,7 @@ class AndroidArgsTest {
 
         flank:
           testShards: 7
+          shardTime: 60
           repeatTests: 8
           files-to-download:
             - /sdcard/screenshots
@@ -158,6 +159,7 @@ class AndroidArgsTest {
 
             // FlankYml
             assert(testShards, 7)
+            assert(shardTime, 60)
             assert(repeatTests, 8)
             assert(filesToDownload, listOf("/sdcard/screenshots", "/sdcard/screenshots2"))
             assert(
@@ -211,6 +213,7 @@ AndroidArgs
 
     flank:
       testShards: 7
+      shardTime: 60
       repeatTests: 8
       smartFlankGcsPath:${' '}
       files-to-download:
@@ -669,6 +672,23 @@ AndroidArgs
         assertThat(AndroidArgs.load(yaml, cli).testShards).isEqualTo(3)
     }
 
+    @Test
+    fun `cli shardTime`() {
+        val cli = AndroidRunCommand()
+        CommandLine(cli).parse("--shard-time=3")
+
+        val yaml = """
+        gcloud:
+          app: $appApk
+          test: $testApk
+
+        flank:
+          shardTime: 2
+      """
+        assertThat(AndroidArgs.load(yaml).shardTime).isEqualTo(2)
+        assertThat(AndroidArgs.load(yaml, cli).shardTime).isEqualTo(3)
+    }
+
     @Test
     fun cli_repeatTests() {
         val cli = AndroidRunCommand()

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

@@ -25,9 +25,10 @@ class FlankYmlTest {
     fun testValidArgs() {
         FlankYml()
         FlankYml(FlankYmlParams(testShards = -1))
-        val yml = FlankYml(FlankYmlParams(testShards = 1, repeatTests = 1))
+        val yml = FlankYml(FlankYmlParams(testShards = 1, repeatTests = 1, shardTime = 58))
         assertThat(yml.flank.repeatTests).isEqualTo(1)
         assertThat(yml.flank.testShards).isEqualTo(1)
+        assertThat(yml.flank.shardTime).isEqualTo(58)
         assertThat(yml.flank.testTargetsAlwaysRun).isEqualTo(emptyList<String>())
         assertThat(FlankYml.map).isNotEmpty()
     }
@@ -38,6 +39,12 @@ class FlankYmlTest {
         FlankYml(FlankYmlParams(testShards = -2))
     }
 
+    @Test
+    fun testInvalidShardTime() {
+        exceptionRule.expectMessage("shardTime must be >= 1 or -1")
+        FlankYml(FlankYmlParams(shardTime = -2))
+    }
+
     @Test
     fun testInvalidrepeatTests() {
         exceptionRule.expectMessage("repeatTests must be >= 1")

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

@@ -53,6 +53,7 @@ class IosArgsTest {
 
         flank:
           testShards: 7
+          shardTime: 60
           repeatTests: 8
           files-to-download:
             - /sdcard/screenshots
@@ -119,6 +120,7 @@ class IosArgsTest {
 
             // FlankYml
             assert(testShards, 7)
+            assert(shardTime, 60)
             assert(repeatTests, 8)
             assert(testTargetsAlwaysRun, listOf("a/testGrantPermissions", "a/testGrantPermissions2"))
 
@@ -160,6 +162,7 @@ IosArgs
 
     flank:
       testShards: 7
+      shardTime: 60
       repeatTests: 8
       smartFlankGcsPath:${' '}
       test-targets-always-run:
@@ -201,6 +204,7 @@ IosArgs
 
             // FlankYml
             assert(testShards, 1)
+            assert(shardTime, -1)
             assert(repeatTests, 1)
             assert(testTargetsAlwaysRun, emptyList<String>())
             assert(filesToDownload, emptyList<String>())
@@ -372,6 +376,23 @@ IosArgs
         assertThat(IosArgs.load(yaml, cli).testShards).isEqualTo(3)
     }
 
+    @Test
+    fun `cli shardTime`() {
+        val cli = IosRunCommand()
+        CommandLine(cli).parse("--shard-time=3")
+
+        val yaml = """
+        gcloud:
+          test: $testPath
+          xctestrun-file: $xctestrunFile
+
+        flank:
+          shardTime: 2
+      """
+        assertThat(IosArgs.load(yaml).shardTime).isEqualTo(2)
+        assertThat(IosArgs.load(yaml, cli).shardTime).isEqualTo(3)
+    }
+
     @Test
     fun cli_repeatTests() {
         val cli = IosRunCommand()

test_runner/src/test/kotlin/ftl/cli/firebase/test/android/AndroidRunCommandTest.kt

@@ -78,6 +78,7 @@ class AndroidRunCommandTest {
         assertThat(cmd.project).isNull()
         assertThat(cmd.resultsHistoryName).isNull()
         assertThat(cmd.testShards).isNull()
+        assertThat(cmd.shardTime).isNull()
         assertThat(cmd.repeatTests).isNull()
         assertThat(cmd.testTargetsAlwaysRun).isNull()
         assertThat(cmd.filesToDownload).isNull()
@@ -291,4 +292,12 @@ class AndroidRunCommandTest {
 
         assertThat(cmd.flakyTestAttempts).isEqualTo(10)
     }
+
+    @Test
+    fun `shardTime parse`() {
+        val cmd = AndroidRunCommand()
+        CommandLine(cmd).parse("--shard-time=99")
+
+        assertThat(cmd.shardTime).isEqualTo(99)
+    }
 }

test_runner/src/test/kotlin/ftl/cli/firebase/test/ios/IosRunCommandTest.kt

@@ -66,6 +66,7 @@ class IosRunCommandTest {
         assertThat(cmd.project).isNull()
         assertThat(cmd.resultsHistoryName).isNull()
         assertThat(cmd.testShards).isNull()
+        assertThat(cmd.shardTime).isNull()
         assertThat(cmd.repeatTests).isNull()
         assertThat(cmd.testTargetsAlwaysRun).isNull()
         assertThat(cmd.testTargets).isNull()
@@ -225,4 +226,12 @@ class IosRunCommandTest {
 
         assertThat(cmd.flakyTestAttempts).isEqualTo(10)
     }
+
+    @Test
+    fun `shardTime parse`() {
+        val cmd = IosRunCommand()
+        CommandLine(cmd).parse("--shard-time=99")
+
+        assertThat(cmd.shardTime).isEqualTo(99)
+    }
 }