* Simplify beforeRunTests return signature

* Treat Args as context

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

@@ -44,10 +44,11 @@ class AndroidRunCommand : CommonRunCommand(), Runnable {
             MockServer.start()
         }
 
-        val config = AndroidArgs.load(Paths.get(configPath), cli = this).validate()
-        runBlocking {
-            if (dumpShards) dumpShards(args = config)
-            else newTestRun(config)
+        AndroidArgs.load(Paths.get(configPath), cli = this).validate().run {
+            runBlocking {
+                if (dumpShards) dumpShards()
+                else newTestRun()
+            }
         }
     }
 

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

@@ -44,12 +44,9 @@ class IosRunCommand : CommonRunCommand(), Runnable {
             MockServer.start()
         }
 
-        val config = IosArgs.load(Paths.get(configPath), cli = this).validate()
-
-        if (dumpShards) {
-            dumpShards(args = config)
-        } else runBlocking {
-            newTestRun(config)
+        IosArgs.load(Paths.get(configPath), cli = this).validate().run {
+            if (dumpShards) dumpShards()
+            else runBlocking { newTestRun() }
         }
     }
 

test_runner/src/main/kotlin/ftl/run/DumpShards.kt

@@ -1,5 +1,6 @@
 package ftl.run
 
+import com.google.common.annotations.VisibleForTesting
 import ftl.args.AndroidArgs
 import ftl.args.IosArgs
 import ftl.args.isInstrumentationTest
@@ -12,31 +13,31 @@ import ftl.util.obfuscatePrettyPrinter
 import java.nio.file.Files
 import java.nio.file.Paths
 
-suspend fun dumpShards(
-    args: AndroidArgs,
+suspend fun AndroidArgs.dumpShards(
+    @VisibleForTesting
     shardFilePath: String = ANDROID_SHARD_FILE,
 ) {
-    if (!args.isInstrumentationTest) throw FlankConfigurationError(
+    if (!isInstrumentationTest) throw FlankConfigurationError(
         "Cannot dump shards for non instrumentation test, ensure test apk has been set."
     )
-    val shards: AndroidMatrixTestShards = args.getAndroidMatrixShards()
+    val shards: AndroidMatrixTestShards = getAndroidMatrixShards()
     saveShardChunks(
         shardFilePath = shardFilePath,
         shards = shards,
         size = shards.flatMap { it.value.shards.values }.count(),
-        obfuscatedOutput = args.obfuscateDumpShards
+        obfuscatedOutput = obfuscateDumpShards
     )
 }
 
-fun dumpShards(
-    args: IosArgs,
+fun IosArgs.dumpShards(
+    @VisibleForTesting
     shardFilePath: String = IOS_SHARD_FILE,
 ) {
     saveShardChunks(
         shardFilePath = shardFilePath,
-        shards = args.testShardChunks.testCases,
-        size = args.testShardChunks.size,
-        obfuscatedOutput = args.obfuscateDumpShards
+        shards = testShardChunks.testCases,
+        size = testShardChunks.size,
+        obfuscatedOutput = obfuscateDumpShards
     )
 }
 

test_runner/src/main/kotlin/ftl/run/NewTestRun.kt

@@ -18,29 +18,30 @@ import ftl.run.platform.runIosTests
 import kotlinx.coroutines.TimeoutCancellationException
 import kotlinx.coroutines.withTimeoutOrNull
 
-suspend fun newTestRun(args: IArgs) = withTimeoutOrNull(args.parsedTimeout) {
+suspend fun IArgs.newTestRun() = withTimeoutOrNull(parsedTimeout) {
+    val args: IArgs = this@newTestRun
     println(args)
-    val (matrixMap, testShardChunks, ignoredTests) = cancelTestsOnTimeout(args.project) { runTests(args) }
+    val (matrixMap, testShardChunks, ignoredTests) =
+        cancelTestsOnTimeout(project) { runTests() }
 
     if (!args.async) {
         cancelTestsOnTimeout(args.project, matrixMap.map) { pollMatrices(matrixMap.map.keys, args).updateMatrixMap(matrixMap) }
         ReportManager.generate(matrixMap, args, testShardChunks, ignoredTests)
         cancelTestsOnTimeout(args.project, matrixMap.map) { fetchArtifacts(matrixMap, args) }
 
         println()
-        matrixMap.printMatricesWebLinks(args.project)
+        matrixMap.printMatricesWebLinks(project)
 
-        matrixMap.validate(args.ignoreFailedTests)
+        matrixMap.validate(ignoreFailedTests)
     }
 }
 
-private suspend fun runTests(args: IArgs): TestResult {
-    return when (args) {
-        is AndroidArgs -> runAndroidTests(args)
-        is IosArgs -> runIosTests(args)
+private suspend fun IArgs.runTests(): TestResult =
+    when (this) {
+        is AndroidArgs -> runAndroidTests()
+        is IosArgs -> runIosTests()
         else -> throw FlankGeneralError("Unknown config type")
     }
-}
 
 private suspend fun <T> cancelTestsOnTimeout(
     projectId: String,

test_runner/src/main/kotlin/ftl/run/RefreshLastRun.kt

@@ -73,7 +73,7 @@ private suspend fun refreshMatrices(matrixMap: MatrixMap, args: IArgs) = corouti
 
     if (dirty) {
         println(FtlConstants.indent + "Updating matrix file")
-        updateMatrixFile(matrixMap, args)
+        args.updateMatrixFile(matrixMap)
     }
     println()
 }

test_runner/src/main/kotlin/ftl/run/common/FetchArtifacts.kt

@@ -57,7 +57,7 @@ internal suspend fun fetchArtifacts(matrixMap: MatrixMap, args: IArgs) = corouti
 
     if (dirty) {
         println(FtlConstants.indent + "Updating matrix file")
-        updateMatrixFile(matrixMap, args)
+        args.updateMatrixFile(matrixMap)
         println()
     }
 }

test_runner/src/main/kotlin/ftl/run/common/UpdateMatrixFile.kt

@@ -7,12 +7,10 @@ import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.Paths
 
-internal fun updateMatrixFile(matrixMap: MatrixMap, args: IArgs): Path {
-    val matrixIdsPath = if (args.useLocalResultDir()) {
-        Paths.get(args.localResultDir, FtlConstants.matrixIdsFile)
-    } else {
-        Paths.get(args.localResultDir, matrixMap.runPath, FtlConstants.matrixIdsFile)
-    }
+internal fun IArgs.updateMatrixFile(matrixMap: MatrixMap): Path {
+    val matrixIdsPath = if (useLocalResultDir())
+        Paths.get(localResultDir, FtlConstants.matrixIdsFile) else
+        Paths.get(localResultDir, matrixMap.runPath, FtlConstants.matrixIdsFile)
     matrixIdsPath.parent.toFile().mkdirs()
     Files.write(matrixIdsPath, prettyPrint.toJson(matrixMap.map).toByteArray())
     return matrixIdsPath

test_runner/src/main/kotlin/ftl/run/platform/RunAndroidTests.kt

@@ -34,32 +34,31 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.coroutineScope
 
-internal suspend fun runAndroidTests(args: AndroidArgs): TestResult = coroutineScope {
-    val (stopwatch, runGcsPath) = beforeRunTests(args)
-
-    // GcAndroidMatrix => GcAndroidTestMatrix
-    // GcAndroidTestMatrix.execute() 3x retry => matrix id (string)
-    val devices = GcAndroidDevice.build(args.devices)
+internal suspend fun AndroidArgs.runAndroidTests(): TestResult = coroutineScope {
+    val args = this@runAndroidTests
+    val stopwatch = beforeRunTests()
+    val devices = GcAndroidDevice.build(devices)
     val testMatrices = mutableListOf<Deferred<TestMatrix>>()
     val allTestShardChunks = mutableListOf<Chunk>()
     val ignoredTestsShardChunks = mutableListOf<List<String>>()
-
     val history = GcToolResults.createToolResultsHistory(args)
-    val otherGcsFiles = args.uploadOtherFiles(runGcsPath)
-    val additionalApks = args.uploadAdditionalApks(runGcsPath)
-    val obbFiles = args.uploadObbFiles(runGcsPath)
+    val otherGcsFiles = uploadOtherFiles()
+    val additionalApks = uploadAdditionalApks()
+    val obbFiles = uploadObbFiles()
 
-    args.createAndroidTestContexts().dumpShards(args).upload(args.resultsBucket, runGcsPath)
+    createAndroidTestContexts()
+        .dumpShards(args)
+        .upload(resultsBucket, resultsDir)
         .forEachIndexed { contextIndex, context ->
             if (context is InstrumentationTestContext) {
                 ignoredTestsShardChunks += context.ignoredTestCases
                 allTestShardChunks += context.shards
             }
-            val androidTestConfig = args.createAndroidTestConfig(context)
-            testMatrices += executeAndroidTestMatrix(runCount = args.repeatTests) { runIndex ->
+            val androidTestConfig = createAndroidTestConfig(context)
+            testMatrices += executeAndroidTestMatrix(runCount = repeatTests) { runIndex ->
                 GcAndroidTestMatrix.build(
                     androidTestConfig = androidTestConfig,
-                    runGcsPath = runGcsPath.createGcsPath(contextIndex, runIndex),
+                    runGcsPath = resultsDir.createGcsPath(contextIndex, runIndex),
                     additionalApkGcsPaths = additionalApks,
                     androidDeviceList = devices,
                     args = args,
@@ -72,10 +71,10 @@ internal suspend fun runAndroidTests(args: AndroidArgs): TestResult = coroutineS
 
     if (testMatrices.isEmpty()) throw FlankGeneralError("There are no tests to run.")
 
-    println(beforeRunMessage(args, allTestShardChunks))
+    println(beforeRunMessage(allTestShardChunks))
 
     TestResult(
-        matrixMap = afterRunTests(testMatrices.awaitAll(), runGcsPath, stopwatch, args),
+        matrixMap = afterRunTests(testMatrices.awaitAll(), stopwatch),
         shardChunks = allTestShardChunks.testCases,
         ignoredTests = ignoredTestsShardChunks.flatten()
     )

test_runner/src/main/kotlin/ftl/run/platform/RunIosTests.kt

@@ -2,7 +2,6 @@ package ftl.run.platform
 
 import com.google.api.services.testing.model.TestMatrix
 import ftl.args.IosArgs
-import ftl.config.FtlConstants
 import ftl.gc.GcIosMatrix
 import ftl.gc.GcIosTestMatrix
 import ftl.gc.GcStorage
@@ -19,6 +18,8 @@ import ftl.run.platform.common.beforeRunMessage
 import ftl.run.platform.common.beforeRunTests
 import ftl.shard.testCases
 import ftl.util.ShardCounter
+import ftl.util.asFileReference
+import ftl.util.uploadIfNeeded
 import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.async
@@ -28,38 +29,39 @@ import kotlinx.coroutines.coroutineScope
 // https://github.com/bootstraponline/gcloud_cli/blob/5bcba57e825fc98e690281cf69484b7ba4eb668a/google-cloud-sdk/lib/googlecloudsdk/api_lib/firebase/test/ios/matrix_creator.py#L109
 // https://cloud.google.com/sdk/gcloud/reference/alpha/firebase/test/ios/run
 // https://cloud.google.com/sdk/gcloud/reference/alpha/firebase/test/ios/
-internal suspend fun runIosTests(iosArgs: IosArgs): TestResult = coroutineScope {
-    val (stopwatch, runGcsPath) = beforeRunTests(iosArgs)
+internal suspend fun IosArgs.runIosTests(): TestResult = coroutineScope {
+    val args = this@runIosTests
+    val stopwatch = beforeRunTests()
 
-    val iosDeviceList = GcIosMatrix.build(iosArgs.devices)
-    val xcTestParsed = parseToNSDictionary(iosArgs.xctestrunFile)
+    val iosDeviceList = GcIosMatrix.build(devices)
+    val xcTestParsed = parseToNSDictionary(xctestrunFile)
 
     val jobs = arrayListOf<Deferred<TestMatrix>>()
-    val runCount = iosArgs.repeatTests
+    val runCount = repeatTests
     val shardCounter = ShardCounter()
-    val history = GcToolResults.createToolResultsHistory(iosArgs)
-    val otherGcsFiles = iosArgs.uploadOtherFiles(runGcsPath)
-    val additionalIpasGcsFiles = iosArgs.uploadAdditionalIpas(runGcsPath)
+    val history = GcToolResults.createToolResultsHistory(args)
+    val otherGcsFiles = uploadOtherFiles()
+    val additionalIpasGcsFiles = uploadAdditionalIpas()
 
-    dumpShards(iosArgs)
-    if (iosArgs.disableResultsUpload.not())
-        GcStorage.upload(IOS_SHARD_FILE, iosArgs.resultsBucket, iosArgs.resultsDir)
+    dumpShards()
+    if (disableResultsUpload.not())
+        GcStorage.upload(IOS_SHARD_FILE, resultsBucket, resultsDir)
 
     // Upload only after parsing shards to detect missing methods early.
-    val xcTestGcsPath = getXcTestGcPath(iosArgs, runGcsPath)
+    val xcTestGcsPath = uploadIfNeeded(xctestrunZip.asFileReference()).gcs
 
-    println(beforeRunMessage(iosArgs, iosArgs.testShardChunks))
+    println(beforeRunMessage(testShardChunks))
 
     repeat(runCount) {
-        jobs += iosArgs.testShardChunks.map { testTargets ->
+        jobs += testShardChunks.map { testTargets ->
             async(Dispatchers.IO) {
                 GcIosTestMatrix.build(
                     iosDeviceList = iosDeviceList,
                     testZipGcsPath = xcTestGcsPath,
-                    runGcsPath = runGcsPath,
+                    runGcsPath = resultsDir,
                     testTargets = testTargets.testMethodNames,
                     xcTestParsed = xcTestParsed,
-                    args = iosArgs,
+                    args = args,
                     shardCounter = shardCounter,
                     toolResultsHistory = history,
                     otherFiles = otherGcsFiles,
@@ -70,16 +72,7 @@ internal suspend fun runIosTests(iosArgs: IosArgs): TestResult = coroutineScope
     }
 
     TestResult(
-        matrixMap = afterRunTests(jobs.awaitAll(), runGcsPath, stopwatch, iosArgs),
-        shardChunks = iosArgs.testShardChunks.testCases
+        matrixMap = afterRunTests(jobs.awaitAll(), stopwatch),
+        shardChunks = testShardChunks.testCases
     )
 }
-
-private fun getXcTestGcPath(
-    iosArgs: IosArgs,
-    runGcsPath: String
-) = if (iosArgs.xctestrunZip.startsWith(FtlConstants.GCS_PREFIX)) {
-    iosArgs.xctestrunZip
-} else {
-    GcStorage.uploadXCTestZip(iosArgs, runGcsPath)
-}

test_runner/src/main/kotlin/ftl/run/platform/android/UploadApks.kt

@@ -47,11 +47,11 @@ private fun GameLoopContext.upload(rootGcsBucket: String, runGcsPath: String) =
 private fun SanityRoboTestContext.upload(rootGcsBucket: String, runGcsPath: String) =
     SanityRoboTestContext(app.uploadIfNeeded(rootGcsBucket, runGcsPath))
 
-suspend fun AndroidArgs.uploadAdditionalApks(runGcsPath: String) =
-    additionalApks.uploadToGcloudIfNeeded(runGcsPath, resultsBucket)
+suspend fun AndroidArgs.uploadAdditionalApks() =
+    additionalApks.uploadToGcloudIfNeeded(resultsDir, resultsBucket)
 
-suspend fun IosArgs.uploadAdditionalIpas(runGcsPath: String) =
-    additionalIpas.uploadToGcloudIfNeeded(runGcsPath, resultsBucket)
+suspend fun IosArgs.uploadAdditionalIpas() =
+    additionalIpas.uploadToGcloudIfNeeded(resultsDir, resultsBucket)
 
 private suspend fun List<String>.uploadToGcloudIfNeeded(
     runGcsPath: String,

test_runner/src/main/kotlin/ftl/run/platform/android/UploadOtherFiles.kt

@@ -8,17 +8,15 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.coroutineScope
 
-internal suspend fun IArgs.uploadOtherFiles(
-    runGcsPath: String
-): Map<String, String> = coroutineScope {
+internal suspend fun IArgs.uploadOtherFiles(): Map<String, String> = coroutineScope {
     otherFiles
         .map { (devicePath: String, filePath: String) ->
-            async(Dispatchers.IO) { devicePath to GcStorage.upload(filePath, resultsBucket, runGcsPath) }
+            async(Dispatchers.IO) { devicePath to GcStorage.upload(filePath, resultsBucket, resultsDir) }
         }.awaitAll().toMap()
 }
 
-internal suspend fun AndroidArgs.uploadObbFiles(runGcsPath: String): Map<String, String> = coroutineScope {
+internal suspend fun AndroidArgs.uploadObbFiles(): Map<String, String> = coroutineScope {
     obbFiles.map {
-        async(Dispatchers.IO) { it to GcStorage.upload(it, resultsBucket, runGcsPath) }
+        async(Dispatchers.IO) { it to GcStorage.upload(it, resultsBucket, resultsDir) }
     }.awaitAll().toMap()
 }

test_runner/src/main/kotlin/ftl/run/platform/common/AfterRunTests.kt

@@ -17,36 +17,34 @@ import kotlinx.coroutines.launch
 import java.nio.file.Files
 import java.nio.file.Paths
 
-internal suspend fun afterRunTests(
+internal suspend fun IArgs.afterRunTests(
     testMatrices: List<TestMatrix>,
-    runGcsPath: String,
     stopwatch: StopWatch,
-    config: IArgs
 ) = MatrixMap(
     map = testMatrices.toSavedMatrixMap(),
-    runPath = runGcsPath
+    runPath = resultsDir
 ).also { matrixMap ->
-    updateMatrixFile(matrixMap, config)
-    saveConfigFile(matrixMap, config)
+    updateMatrixFile(matrixMap)
+    saveConfigFile(matrixMap)
 
     println(FtlConstants.indent + "${matrixMap.map.size} matrix ids created in ${stopwatch.check()}")
     val gcsBucket = "https://console.developers.google.com/storage/browser/" +
-            config.resultsBucket + "/" + matrixMap.runPath
+            resultsBucket + "/" + matrixMap.runPath
     println(FtlConstants.indent + gcsBucket)
     println()
-    matrixMap.printMatricesWebLinks(config.project)
+    matrixMap.printMatricesWebLinks(project)
 }
 
 private fun List<TestMatrix>.toSavedMatrixMap() =
     associate { matrix -> matrix.testMatrixId to createSavedMatrix(matrix) }
 
-private fun saveConfigFile(matrixMap: MatrixMap, args: IArgs) {
-    val configFilePath = if (args.useLocalResultDir())
-        Paths.get(args.localResultDir, FtlConstants.configFileName(args)) else
-        Paths.get(args.localResultDir, matrixMap.runPath, FtlConstants.configFileName(args))
+private fun IArgs.saveConfigFile(matrixMap: MatrixMap) {
+    val configFilePath = if (useLocalResultDir())
+        Paths.get(localResultDir, FtlConstants.configFileName(this)) else
+        Paths.get(localResultDir, matrixMap.runPath, FtlConstants.configFileName(this))
 
     configFilePath.parent.toFile().mkdirs()
-    Files.write(configFilePath, args.data.toByteArray())
+    Files.write(configFilePath, data.toByteArray())
 }
 
 internal suspend inline fun MatrixMap.printMatricesWebLinks(project: String) = coroutineScope {