Create testMatrix to rerun failed tests in FTL

AndroidXCI/lib/src/main/kotlin/dev/androidx/ci/testRunner/FirebaseTestLabController.kt

@@ -111,7 +111,7 @@ internal class FirebaseTestLabController(
         pullScreenshots: Boolean = false,
         cachedTestMatrixFilter: CachedTestMatrixFilter,
         testTargets: List<String>? = null,
-        flakyTestAttempts: Int? = 2
+        flakyTestAttempts: Int = 2
     ): List<TestMatrix> {
         val devices = (devicePicker ?: defaultDevicePicker).pickDevices()
         logger.info {
@@ -134,6 +134,22 @@ internal class FirebaseTestLabController(
         }
     }
 
+    /**
+     * Enqueues a request to create a [TestMatrix] to run the tests
+     * specified in the [testTargets] list using the same configuration as [testMatrix]
+     */
+    suspend fun scheduleTests(
+        testMatrix: TestMatrix,
+        testTargets: List<String>? = null,
+        cachedTestMatrixFilter: CachedTestMatrixFilter
+    ): TestMatrix {
+        return testMatrixStore.getOrCreateTestMatrix(
+            testMatrix = testMatrix,
+            testTargets = testTargets,
+            cachedTestMatrixFilter = cachedTestMatrixFilter
+        )
+    }
+
     suspend fun collectTestResultsByTestMatrixIds(
         testMatrixIds: List<String>,
         pollIntervalMs: Long

AndroidXCI/lib/src/main/kotlin/dev/androidx/ci/testRunner/TestMatrixStore.kt

@@ -49,7 +49,7 @@ internal class TestMatrixStore(
     private val datastoreApi: DatastoreApi,
     private val firebaseTestLabApi: FirebaseTestLabApi,
     toolsResultApi: ToolsResultApi,
-    private val resultsGcsPrefix: GcsPath,
+    private val resultsGcsPrefix: GcsPath
 ) {
     private val logger = logger()
     private val toolsResultStore = ToolsResultStore(
@@ -70,7 +70,7 @@ internal class TestMatrixStore(
         pullScreenshots: Boolean = false,
         cachedTestMatrixFilter: CachedTestMatrixFilter = { true },
         testTargets: List<String>? = null,
-        flakyTestAttempts: Int? = 2
+        flakyTestAttempts: Int = 2
     ): TestMatrix {
 
         val testRunId = TestRun.createId(
@@ -105,21 +105,89 @@ internal class TestMatrixStore(
         logger.trace {
             "No test run history for $testRunId or cached TestMatrix is rejected, creating a new one."
         }
-        val newTestMatrix = firebaseTestLabApi.createTestMatrix(
-            projectId = firebaseProjectId,
-            requestId = UUID.randomUUID().toString(),
-            testMatrix = createNewTestMatrix(
-                testRunKey = testRunId,
-                environmentMatrix = environmentMatrix,
-                clientInfo = clientInfo,
-                sharding = sharding,
-                deviceSetup = deviceSetup,
-                appApk = appApk,
-                testApk = testApk,
-                pullScreenshots = pullScreenshots,
-                testTargets = testTargets,
-                flakyTestAttempts = flakyTestAttempts
-            )
+
+        val newTestMatrix = createNewTestMatrix(
+            testRunKey = testRunId,
+            environmentMatrix = environmentMatrix,
+            clientInfo = clientInfo,
+            sharding = sharding,
+            deviceSetup = deviceSetup,
+            appApk = appApk,
+            testApk = testApk,
+            pullScreenshots = pullScreenshots,
+            testTargets = testTargets,
+            flakyTestAttempts = flakyTestAttempts
+        )
+        logger.info {
+            "created test matrix: $newTestMatrix"
+        }
+        // save it to cache, we don't worry about races here much such that if another instance happens to be creating
+        // the exact same test, one of them will win but that is OK since they'll each use their own test matrices and
+        // followup calls will re-use the winner of this race.
+        val testRun = TestRun(
+            id = testRunId,
+            testMatrixId = checkNotNull(newTestMatrix.testMatrixId) {
+                "newly created test matrix should not have null id $newTestMatrix"
+            }
+        )
+        datastoreApi.put(testRun.toEntity())
+        logger.info {
+            "saved test matrix info: $testRun"
+        }
+        return newTestMatrix
+    }
+
+    /**
+     * Creates a TestMatrix for the given configuration or returns an existing one if we've run the same tests
+     * specified in [testTargets] with the same environment configuration.
+     */
+    suspend fun getOrCreateTestMatrix(
+        testMatrix: TestMatrix,
+        cachedTestMatrixFilter: CachedTestMatrixFilter = { true },
+        testTargets: List<String>? = null,
+        flakyTestAttempts: Int = 2
+    ): TestMatrix {
+        checkNotNull(testMatrix.testMatrixId) {
+            "Test matrix id for the base test matrix should not be null"
+        }
+        val testRunId = TestRun.createId(
+            datastoreApi = datastoreApi,
+            environment = testMatrix.environmentMatrix,
+            clientInfo = testMatrix.clientInfo,
+            sharding = testMatrix.testSpecification.androidInstrumentationTest?.shardingOption,
+            testSetup = testMatrix.testSpecification.testSetup,
+            testTargets = testTargets,
+            testMatrixId = testMatrix.testMatrixId
+        )
+        logger.trace {
+            "test run id: $testRunId"
+        }
+
+        getExistingTestMatrix(testRunId)?.let {
+            logger.info("found existing test matrix: ${it.testMatrixId} with state: ${it.state}")
+            val state = it.state
+            // these states are ordered so anything above ERROR is not worth re-using
+            if (state != null && state >= TestMatrix.State.ERROR) {
+                logger.warn {
+                    "Skipping cache for ${it.testMatrixId} because its state is $state"
+                }
+            } else if (!cachedTestMatrixFilter(it)) {
+                logger.info {
+                    "Not re-using cached matrix due to filter"
+                }
+            } else {
+                return it
+            }
+        }
+        logger.trace {
+            "No test run history for $testRunId or cached TestMatrix is rejected, creating a new one."
+        }
+
+        val newTestMatrix = createNewTestMatrix(
+            testRunKey = testRunId,
+            testMatrix = testMatrix,
+            testTargets = testTargets,
+            flakyTestAttempts = flakyTestAttempts
         )
         logger.info {
             "created test matrix: $newTestMatrix"
@@ -183,7 +251,7 @@ internal class TestMatrixStore(
         testApk: UploadedApk,
         pullScreenshots: Boolean = false,
         testTargets: List<String>? = null,
-        flakyTestAttempts: Int? = 2
+        flakyTestAttempts: Int = 2
     ): TestMatrix {
         val packageName = firebaseTestLabApi.getApkDetails(
             FileReference(
@@ -210,35 +278,96 @@ internal class TestMatrixStore(
                 else null
             )
         }
-        return TestMatrix(
-            projectId = firebaseProjectId,
-            flakyTestAttempts = flakyTestAttempts,
-            testSpecification = TestSpecification(
-                testTimeout = "2700s", // Limit for physical devices.
-                disableVideoRecording = false,
-                disablePerformanceMetrics = true, // Not a useful feature for androidx
-                androidInstrumentationTest = AndroidInstrumentationTest(
-                    appApk = FileReference(
-                        gcsPath = appApk.gcsPath.path
-                    ),
-                    testApk = FileReference(
-                        gcsPath = testApk.gcsPath.path
-                    ),
-                    shardingOption = sharding,
-                    testTargets = testTargets
+        val testSpecification = TestSpecification(
+            testTimeout = "2700s", // Limit for physical devices.
+            disableVideoRecording = false,
+            disablePerformanceMetrics = true, // Not a useful feature for androidx
+            androidInstrumentationTest = AndroidInstrumentationTest(
+                appApk = FileReference(
+                    gcsPath = appApk.gcsPath.path
                 ),
-                testSetup = testSetup
+                testApk = FileReference(
+                    gcsPath = testApk.gcsPath.path
+                ),
+                shardingOption = sharding,
+                testTargets = testTargets
+            ),
+            testSetup = testSetup
+        )
+        val resultStorage = ResultStorage(
+            googleCloudStorage = GoogleCloudStorage(
+                gcsPath = testRunKey.resultGcsPath().path
             ),
+            toolResultsHistory = ToolResultsHistory(
+                projectId = firebaseProjectId,
+                historyId = historyId
+            )
+        )
+        return createTestMatrix(
+            flakyTestAttempts = flakyTestAttempts,
+            testSpecification = testSpecification,
             clientInfo = clientInfo,
             environmentMatrix = environmentMatrix,
-            resultStorage = ResultStorage(
-                googleCloudStorage = GoogleCloudStorage(
-                    gcsPath = testRunKey.resultGcsPath().path
-                ),
-                toolResultsHistory = ToolResultsHistory(
-                    projectId = firebaseProjectId,
-                    historyId = historyId
-                )
+            resultStorage = resultStorage
+        )
+    }
+
+    /**
+     * Creates a [TestMatrix] to run the tests specified in [testTargets] list
+     * using the same configuration as [testMatrix]
+     */
+    private suspend fun createNewTestMatrix(
+        testRunKey: TestRun.Id,
+        testMatrix: TestMatrix,
+        testTargets: List<String>? = null,
+        flakyTestAttempts: Int = 0
+    ): TestMatrix {
+        logger.trace {
+            "test matrix id: ${testMatrix.testMatrixId}"
+        }
+
+        val testSpecification = testMatrix.testSpecification.copy(
+            androidInstrumentationTest = testMatrix.testSpecification.androidInstrumentationTest?.copy(
+                testTargets = testTargets
+            )
+        )
+
+        val resultStorage = ResultStorage(
+            googleCloudStorage = GoogleCloudStorage(
+                gcsPath = testRunKey.resultGcsPath().path
+            ),
+            toolResultsHistory = testMatrix.resultStorage.toolResultsHistory
+        )
+
+        return createTestMatrix(
+            clientInfo = testMatrix.clientInfo,
+            environmentMatrix = testMatrix.environmentMatrix,
+            testSpecification = testSpecification,
+            resultStorage = resultStorage,
+            flakyTestAttempts = flakyTestAttempts
+        )
+    }
+
+    /**
+     * Creates a [TestMatrix] using the specified parameters
+     */
+    suspend fun createTestMatrix(
+        clientInfo: ClientInfo?,
+        environmentMatrix: EnvironmentMatrix,
+        testSpecification: TestSpecification,
+        resultStorage: ResultStorage,
+        flakyTestAttempts: Int
+    ): TestMatrix {
+        return firebaseTestLabApi.createTestMatrix(
+            projectId = firebaseProjectId,
+            requestId = UUID.randomUUID().toString(),
+            testMatrix = TestMatrix(
+                projectId = firebaseProjectId,
+                flakyTestAttempts = flakyTestAttempts,
+                testSpecification = testSpecification,
+                clientInfo = clientInfo,
+                environmentMatrix = environmentMatrix,
+                resultStorage = resultStorage
             )
         )
     }

AndroidXCI/lib/src/main/kotlin/dev/androidx/ci/testRunner/TestRunnerService.kt

@@ -81,9 +81,19 @@ interface TestRunnerService {
         pullScreenshots: Boolean = false,
         cachedTestMatrixFilter: CachedTestMatrixFilter = { true },
         testTargets: List<String>? = null,
-        flakyTestAttempts: Int? = 2
+        flakyTestAttempts: Int = 2
     ): ScheduleTestsResponse
 
+    /**
+     * Schedules a task to create a [TestMatrix] to run tests
+     * specified in the [testTargets] list using the same configuration as [testMatrix]
+     */
+    suspend fun scheduleTests(
+        testMatrix: TestMatrix,
+        testTargets: List<String>? = null,
+        cachedTestMatrixFilter: CachedTestMatrixFilter = { true }
+    ): TestMatrix
+
     /**
      * Queries the Firebase for the given [testMatrixId] and returns it if it exists.
      */

AndroidXCI/lib/src/main/kotlin/dev/androidx/ci/testRunner/TestRunnerServiceImpl.kt

@@ -101,7 +101,7 @@ internal class TestRunnerServiceImpl internal constructor(
         pullScreenshots: Boolean,
         cachedTestMatrixFilter: CachedTestMatrixFilter,
         testTargets: List<String>?,
-        flakyTestAttempts: Int?
+        flakyTestAttempts: Int
     ): TestRunnerService.ScheduleTestsResponse {
         val testMatrices = testLabController.submitTests(
             appApk = appApk ?: apkStore.getPlaceholderApk(),
@@ -120,6 +120,18 @@ internal class TestRunnerServiceImpl internal constructor(
         )
     }
 
+    override suspend fun scheduleTests(
+        testMatrix: TestMatrix,
+        testTargets: List<String>?,
+        cachedTestMatrixFilter: CachedTestMatrixFilter
+    ): TestMatrix {
+        return testLabController.scheduleTests(
+            testMatrix = testMatrix,
+            testTargets = testTargets,
+            cachedTestMatrixFilter = cachedTestMatrixFilter
+        )
+    }
+
     override suspend fun getTestMatrix(
         testMatrixId: String
     ): TestMatrix? {

AndroidXCI/lib/src/main/kotlin/dev/androidx/ci/testRunner/dto/TestRun.kt

@@ -26,6 +26,7 @@ import dev.androidx.ci.datastore.DatastoreApi
 import dev.androidx.ci.generated.ftl.ClientInfo
 import dev.androidx.ci.generated.ftl.EnvironmentMatrix
 import dev.androidx.ci.generated.ftl.ShardingOption
+import dev.androidx.ci.generated.ftl.TestSetup
 import dev.androidx.ci.testRunner.dto.TestRun.Companion.createId
 import dev.androidx.ci.testRunner.vo.ApkInfo
 import dev.androidx.ci.testRunner.vo.DeviceSetup
@@ -92,6 +93,38 @@ internal class TestRun(
             val sha = sha256(json.toByteArray(Charsets.UTF_8))
             return Id(datastoreApi.createKey(datastoreApi.testRunObjectKind, sha))
         }
+
+        /**
+         * Creates a unique ID for the given parameters
+         */
+        fun createId(
+            datastoreApi: DatastoreApi,
+            environment: EnvironmentMatrix,
+            clientInfo: ClientInfo?,
+            testSetup: TestSetup?,
+            sharding: ShardingOption?,
+            testTargets: List<String>?,
+            testMatrixId: String
+        ): Id {
+            val json = adapter.toJson(
+                mapOf(
+                    "e" to environment,
+                    "clientInfo" to clientInfo,
+                    "sharding" to sharding,
+                    "instrumentationArgs" to testSetup?.environmentVariables?.flatMap {
+                        listOf(it.key, it.value)
+                    },
+                    "additionalApks" to testSetup?.additionalApks?.map {
+                        it.location?.gcsPath
+                    }, // The order we install additional apks is important, so we do not sort here.
+                    "directoriesToPull" to testSetup?.directoriesToPull?.sorted(),
+                    "testTargets" to testTargets?.sorted(),
+                    "testMatrixId" to testMatrixId
+                )
+            )
+            val sha = sha256(json.toByteArray(Charsets.UTF_8))
+            return Id(datastoreApi.createKey(datastoreApi.testRunObjectKind, sha))
+        }
     }
 }