add sharding

Author: Piotr Mądry

plugin/build.gradle

@@ -6,7 +6,7 @@ apply plugin: "com.gradle.plugin-publish"
 apply plugin: "org.gradle.kotlin.kotlin-dsl"
 
 group = "firebase.test.lab"
-version = "1.1.6"
+version = "1.1.6.23"
 
 gradlePlugin {
     plugins {

plugin/src/main/java/com/appunite/firebasetestlabplugin/FirebaseTestLabPlugin.kt

@@ -6,9 +6,11 @@ import com.android.build.gradle.api.BaseVariantOutput
 import com.android.build.gradle.api.TestVariant
 import com.appunite.firebasetestlabplugin.cloud.CloudTestResultDownloader
 import com.appunite.firebasetestlabplugin.cloud.FirebaseTestLabProcessCreator
+import com.appunite.firebasetestlabplugin.cloud.ProcessData
 import com.appunite.firebasetestlabplugin.cloud.TestType
 import com.appunite.firebasetestlabplugin.model.Device
 import com.appunite.firebasetestlabplugin.model.TestResults
+import com.appunite.firebasetestlabplugin.tasks.InstrumentationShardingTask
 import com.appunite.firebasetestlabplugin.utils.Constants
 import org.apache.tools.ant.taskdefs.condition.Os
 import org.gradle.api.GradleException
@@ -20,8 +22,9 @@ import org.gradle.kotlin.dsl.closureOf
 import org.gradle.kotlin.dsl.register
 import java.io.ByteArrayOutputStream
 import java.io.File
+import java.io.Serializable
 
-internal class FirebaseTestLabPlugin : Plugin<Project> {
+class FirebaseTestLabPlugin : Plugin<Project> {
 
     open class HiddenExec : Exec() {
         init {
@@ -66,7 +69,7 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
         }
     }
 
-    data class Sdk(val gcloud: File, val gsutil: File)
+    data class Sdk(val gcloud: File, val gsutil: File): Serializable
 
     private fun createDownloadSdkTask(project: Project, cloudSdkPath: String?): Sdk =
             if (cloudSdkPath != null) {
@@ -187,16 +190,9 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
                 throw IllegalStateException("If you want to clear directory before run you need to setup cloudBucketName and cloudDirectoryName")
             }
 
-            val firebaseTestLabProcessCreator = FirebaseTestLabProcessCreator(
-                    sdk,
-                    cloudBucketName,
-                    cloudDirectoryName,
-                    project.logger
-            )
-
             (project.extensions.findByName(ANDROID) as AppExtension).apply {
                 testVariants.toList().forEach { testVariant ->
-                    createGroupedTestLabTask(devices, testVariant, firebaseTestLabProcessCreator, ignoreFailures, downloader)
+                    createGroupedTestLabTask(devices, testVariant, ignoreFailures, downloader, sdk, cloudBucketName, cloudDirectoryName)
                 }
             }
 
@@ -206,14 +202,17 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
 
     data class DeviceAppMap(val device: Device, val apk: BaseVariantOutput)
 
-    data class Test(val device: Device, val apk: BaseVariantOutput, val testApk: BaseVariantOutput)
-
+    data class Test(val device: Device, val apk: BaseVariantOutput, val testApk: BaseVariantOutput): Serializable
+    
     private fun createGroupedTestLabTask(
-            devices: List<Device>,
-            variant: TestVariant,
-            firebaseTestLabProcessCreator: FirebaseTestLabProcessCreator,
-            ignoreFailures: Boolean,
-            downloader: CloudTestResultDownloader?) {
+        devices: List<Device>,
+        variant: TestVariant,
+        ignoreFailures: Boolean,
+        downloader: CloudTestResultDownloader?,
+        sdk: Sdk,
+        cloudBucketName: String?,
+        cloudDirectoryName: String?
+    ) {
         val variantName = variant.testedVariant?.name?.capitalize() ?: ""
 
         val cleanTask = "firebaseTestLabClean${variantName.capitalize()}"
@@ -264,19 +263,58 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
                         dependsOn(taskSetup)
                         dependsOn(arrayOf(test.apk.assemble))
                         doLast {
-                            val result = firebaseTestLabProcessCreator.callFirebaseTestLab(test.device, test.apk.outputFile, TestType.Robo)
+                            val result = FirebaseTestLabProcessCreator.callFirebaseTestLab(ProcessData(
+                                sdk = sdk,
+                                gCloudBucketName = cloudBucketName,
+                                gCloudDirectory = cloudDirectoryName,
+                                device = test.device,
+                                apk = test.apk.outputFile,
+                                
+                                testType = TestType.Robo
+                            ))
                             processResult(result, ignoreFailures)
                         }
                     })
                 }
-
-        val instrumentationTasks: List<Task> = combineAll(appVersions, variant.outputs, {deviceAndMap, testApk -> Test(deviceAndMap.device, deviceAndMap.apk, testApk)})
-                .map {
-                    test ->
-                    val devicePart = test.device.name.capitalize()
-                    val apkPart = dashToCamelCase(test.apk.name).capitalize()
-                    val testApkPart = test.testApk.let { if (it.filters.isEmpty()) "" else dashToCamelCase(it.name).capitalize() }
-                    val taskName = "$runTestsTaskInstrumentation$devicePart$apkPart$testApkPart"
+    
+        val instrumentationTasks: List<Task> = combineAll(appVersions, variant.outputs)
+        { deviceAndMap, testApk -> Test(deviceAndMap.device, deviceAndMap.apk, testApk) }
+            .map { test ->
+                val devicePart = test.device.name.capitalize()
+                val apkPart = dashToCamelCase(test.apk.name).capitalize()
+                val testApkPart = test.testApk.let { if (it.filters.isEmpty()) "" else dashToCamelCase(it.name).capitalize() }
+                val taskName = "$runTestsTaskInstrumentation$devicePart$apkPart$testApkPart"
+                val numShards = test.device.numShards
+    
+                val file = File(project.buildDir, "TestResults.txt")
+    
+                if (numShards > 0) {
+                    project.tasks.create(taskName, InstrumentationShardingTask::class.java) {
+                        group = Constants.FIREBASE_TEST_LAB
+                        description = "Run Instrumentation test for ${test.device.name} device on $variantName/${test.apk.name} in Firebase Test Lab"
+                        this.processData = ProcessData(
+                            sdk = sdk,
+                            gCloudBucketName = cloudBucketName,
+                            gCloudDirectory = cloudDirectoryName,
+                            device = test.device,
+                            apk = test.apk.outputFile,
+                            testType = TestType.Instrumentation(test.testApk.outputFile)
+                        )
+                        this.stateFile = file
+            
+                        if (downloader != null) {
+                            mustRunAfter(cleanTask)
+                        }
+                        dependsOn(taskSetup)
+                        dependsOn(arrayOf(test.apk.assemble, test.testApk.assemble))
+            
+                        doLast {
+                            val resultCode = file.readText().toInt()
+                            processResult(resultCode, ignoreFailures)
+                        }
+                    }
+        
+                } else {
                     project.task(taskName, closureOf<Task> {
                         inputs.files(test.testApk.outputFile, test.apk.outputFile)
                         group = Constants.FIREBASE_TEST_LAB
@@ -287,11 +325,19 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
                         dependsOn(taskSetup)
                         dependsOn(arrayOf(test.apk.assemble, test.testApk.assemble))
                         doLast {
-                            val result = firebaseTestLabProcessCreator.callFirebaseTestLab(test.device, test.apk.outputFile, TestType.Instrumentation(test.testApk.outputFile))
+                            val result = FirebaseTestLabProcessCreator.callFirebaseTestLab(ProcessData(
+                                sdk = sdk,
+                                gCloudBucketName = cloudBucketName,
+                                gCloudDirectory = cloudDirectoryName,
+                                device = test.device,
+                                apk = test.apk.outputFile,
+                                testType = TestType.Instrumentation(test.testApk.outputFile)
+                            ))
                             processResult(result, ignoreFailures)
                         }
                     })
                 }
+            }
 
         val allInstrumentation: Task = project.task(runTestsTaskInstrumentation, closureOf<Task> {
             group = Constants.FIREBASE_TEST_LAB
@@ -334,63 +380,13 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
                 if (roboTasks.isEmpty()) throw IllegalStateException("Nothing match your filter")
             }
         })
-    
-        combineAll(appVersions, variant.outputs) { deviceAndMap, testApk -> Test(deviceAndMap.device, deviceAndMap.apk, testApk) }
-            .map { test ->
-                val devicePart = test.device.name.capitalize()
-                val apkPart = dashToCamelCase(test.apk.name).capitalize()
-                val testApkPart = test.testApk.let { if (it.filters.isEmpty()) "" else dashToCamelCase(it.name).capitalize() }
-                val numShards = test.device.numShards
-                if (numShards > 1) {
-                    
-                    val shardedTasks: List<Task> = (0 until test.device.numShards).map { shardIndex ->
-                        val taskName = "$runTestsTaskInstrumentation$devicePart$apkPart$testApkPart" + "NumShards$numShards" + "ShardIndex$shardIndex"
-                        project.task(taskName, closureOf<Task> {
-                            inputs.files(test.testApk.outputFile, test.apk.outputFile)
-                            group = Constants.FIREBASE_TEST_LAB
-                            description = "Run Instrumentation test for ${test.device.name} device on $variantName/${test.apk.name} with ShardNumber $numShards and ShardIndex $shardIndex in Firebase Test Lab"
-                            if (downloader != null) {
-                                mustRunAfter(cleanTask)
-                            }
-                            dependsOn(taskSetup)
-                            dependsOn(arrayOf(test.apk.assemble, test.testApk.assemble))
-                            doLast {
-                                val result = firebaseTestLabProcessCreator.callFirebaseTestLab(test.device, test.apk.outputFile, TestType.Instrumentation(test.testApk.outputFile), shardIndex)
-                                processResult(result, ignoreFailures)
-                            }
-                        })
-                    }
-    
-                    val runTestsTaskSharded = "firebaseTestLabExecuteAllShardForConfiguration$devicePart$apkPart$testApkPart"
-                    project.task(runTestsTaskSharded, closureOf<Task> {
-                        group = Constants.FIREBASE_TEST_LAB
-                        description = "Run all tests for $variantName sharded in Firebase Test Lab"
-                        dependsOn(shardedTasks)
-        
-                        doFirst {
-                            if (devices.isEmpty()) throw IllegalStateException("You need to set et least one device in:\n" +
-                                "firebaseTestLab {" +
-                                "  devices {\n" +
-                                "    nexus6 {\n" +
-                                "      androidApiLevels = [21]\n" +
-                                "      deviceIds = [\"Nexus6\"]\n" +
-                                "      locales = [\"en\"]\n" +
-                                "    }\n" +
-                                "  } " +
-                                "}")
-            
-                            if (roboTasks.isEmpty()) throw IllegalStateException("Nothing match your filter")
-                        }
-                    })
-                }
-            }
 
         project.task(runTestsTask, closureOf<Task> {
             group = Constants.FIREBASE_TEST_LAB
             description = "Run all tests for $variantName in Firebase Test Lab"
             dependsOn(allRobo, allInstrumentation)
         })
-
+    
         if (downloader != null) {
             project.task(downloadTask, closureOf<Task> {
                 group = Constants.FIREBASE_TEST_LAB
@@ -405,7 +401,7 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
             })
         }
     }
-
+    
     private fun processResult(result: TestResults, ignoreFailures: Boolean) {
         if (result.isSuccessful) {
             project.logger.lifecycle(result.message)
@@ -417,11 +413,21 @@ internal class FirebaseTestLabPlugin : Plugin<Project> {
             }
         }
     }
+    
+    private fun processResult(resultCode: Int, ignoreFailures: Boolean) =
+        if (resultCode == 0) {
+            project.logger.lifecycle("SUCCESS: All tests passed.")
+        } else {
+            if (ignoreFailures) {
+                project.logger.error("FAILURE: Tests failed.")
+            } else {
+                throw GradleException("FAILURE: Tests failed.")
+            }
+        }
 }
 
-
 private fun <T1, T2, R> combineAll(l1: Collection<T1>, l2: Collection<T2>, func: (T1, T2) -> R): List<R> =
         l1.flatMap { t1 -> l2.map { t2 -> func(t1, t2)} }
 
-private fun dashToCamelCase(dash: String): String =
+fun dashToCamelCase(dash: String): String =
         dash.split('-', '_').joinToString("") { it.capitalize() }

plugin/src/main/java/com/appunite/firebasetestlabplugin/cloud/FirebaseTestLabProcessCreator.kt

@@ -4,19 +4,25 @@ import com.appunite.firebasetestlabplugin.FirebaseTestLabPlugin
 import com.appunite.firebasetestlabplugin.model.Device
 import com.appunite.firebasetestlabplugin.model.TestResults
 import com.appunite.firebasetestlabplugin.utils.joinArgs
-import org.gradle.api.logging.Logger
 import java.io.File
+import java.io.Serializable
 
-sealed class TestType {
+sealed class TestType : Serializable {
     object Robo : TestType()
     data class Instrumentation(val testApk: File) : TestType()
 }
 
-internal class FirebaseTestLabProcessCreator(
-        private val sdk: FirebaseTestLabPlugin.Sdk,
-        private val gCloudBucketName: String?,
-        private val gCloudDirectory: String?,
-        private val logger: Logger) {
+data class ProcessData(
+    val sdk: FirebaseTestLabPlugin.Sdk,
+    val gCloudBucketName: String?,
+    val gCloudDirectory: String?,
+    val device: Device,
+    val apk: File,
+    val testType: TestType,
+    val shardIndex: Int = -1
+) : Serializable
+
+object FirebaseTestLabProcessCreator {
 
     private val resultMessageMap = mapOf(
             0 to "All test executions passed.",
@@ -29,39 +35,39 @@ internal class FirebaseTestLabProcessCreator(
             20 to "A test infrastructure error occurred."
     )
 
-    fun callFirebaseTestLab(device: Device, apk: File, testType: TestType, shardIndex: Int = -1): TestResults {
-        val processBuilder = createProcess(device, apk, testType)
-        logger.debug(processBuilder.command().joinToString(separator = " ", transform = { "\"$it\"" }))
-        val process = processBuilder.start()
-        process.errorStream.bufferedReader().forEachLine { logger.lifecycle(it) }
-        process.inputStream.bufferedReader().forEachLine { logger.lifecycle(it) }
-
-        val resultCode = process.waitFor()
+    fun callFirebaseTestLab(processData: ProcessData): TestResults {
+        val process: Process = createProcess(processData).start()
+        val resultCode: Int = process.let {
+            it.errorStream.bufferedReader().forEachLine { errorInfo -> println(errorInfo) }
+            it.inputStream.bufferedReader().forEachLine { info -> println(info) }
+            process.waitFor()
+        }
         return TestResults(
-                isSuccessful = resultCode == 0,
-                message = resultMessageMap[resultCode] ?: ""
+            isSuccessful = resultCode == 0,
+            message = resultMessageMap.getOrElse(resultCode) { "Unknown error with code: $resultCode" }
         )
     }
 
-    private fun createProcess(device: Device, apk: File, testType: TestType, shardIndex: Int = -1): ProcessBuilder {
+    private fun createProcess(processData: ProcessData): ProcessBuilder {
+        val device: Device = processData.device
         return ProcessBuilder(
             sequenceOf(
-                sdk.gcloud.absolutePath,
+                processData.sdk.gcloud.absolutePath,
                 "firebase", "test", "android", "run",
                 "--format=json",
                 "--device-ids=${device.deviceIds.joinArgs()}",
-                "--app=$apk",
+                "--app=$processData.apk",
                 "--locales=${device.locales.joinArgs()}",
                 "--os-version-ids=${device.androidApiLevels.joinArgs()}",
                 "--orientations=${device.screenOrientations.map { orientation -> orientation.gcloudName }.joinArgs()}")
-                .plus(when (testType) {
+                .plus(when (processData.testType) {
                     TestType.Robo -> sequenceOf("--type=robo")
-                    is TestType.Instrumentation -> sequenceOf("--type=instrumentation", "--test=${testType.testApk}")
+                    is TestType.Instrumentation -> sequenceOf("--type=instrumentation", "--test=${processData.testType.testApk}")
                 })
-                .plus(gCloudBucketName?.let { sequenceOf("--results-bucket=$it") } ?: sequenceOf())
-                .plus(gCloudDirectory?.let { sequenceOf("--results-dir=$it") } ?: sequenceOf())
+                .plus(processData.gCloudBucketName?.let { sequenceOf("--results-bucket=$it") } ?: sequenceOf())
+                .plus(processData.gCloudDirectory?.let { sequenceOf("--results-dir=$it") } ?: sequenceOf())
                 .plus(if (device.isUseOrchestrator) sequenceOf("--use-orchestrator") else sequenceOf())
-                .plus(if (device.environmentVariables.isNotEmpty()) sequenceOf("--environment-variables=${device.environmentVariables.joinToString(",")}${addSharding(device, shardIndex)}") else sequenceOf())
+                .plus(setupEnvironmentVariables(device, processData.shardIndex))
                 .plus(if (device.testTargets.isNotEmpty()) sequenceOf("--test-targets=${device.testTargets.joinToString(",")}") else sequenceOf())
                 .plus(device.customParamsForGCloudTool)
                 .plus(device.testRunnerClass?.let { sequenceOf("--test-runner-class=$it") } ?: sequenceOf())
@@ -70,8 +76,13 @@ internal class FirebaseTestLabProcessCreator(
         )
     }
 
-    private fun addSharding(device: Device, shardIndex: Int): String = when {
-        device.numShards <= 0 && shardIndex >= 0 -> ", numShards=${device.numShards}, shardIndex=$shardIndex"
-        else -> ""
-    }
+    private fun setupEnvironmentVariables(device: Device, shardIndex: Int): Sequence<String> =
+        if (device.environmentVariables.isNotEmpty() || device.numShards > 0)
+            sequenceOf(StringBuilder()
+                .append("--environment-variables=")
+                .append(if (device.environmentVariables.isNotEmpty()) device.environmentVariables.joinToString(",") else "")
+                .append(if (device.environmentVariables.isNotEmpty() && device.numShards > 0) "," else "")
+                .append(if (device.numShards > 0) "numShards=${device.numShards},shardIndex=$shardIndex" else "")
+                .toString())
+        else sequenceOf()
 }

plugin/src/main/java/com/appunite/firebasetestlabplugin/model/Device.kt

@@ -1,6 +1,8 @@
 package com.appunite.firebasetestlabplugin.model
 
-class Device(val name: String) {
+import java.io.Serializable
+
+class Device(val name: String) : Serializable {
     var locales: List<String> = listOf("en")
     var screenOrientations: List<ScreenOrientation> = listOf(ScreenOrientation.PORTRAIT)
     var androidApiLevels: List<Int> = listOf()