Allow overriding the on-device coverage agent version

The agent is released in lockstep with AGP and resolves at the plugin
version by default. A project that hits an agent bug fixed in another
release can now pin that release without changing its AGP version:

  android.experimental.androidTest.coverage.onDeviceAgentVersion=<version>

Bug: 571657536
Test: CoverageAgentConfigurationsTest. JacocoTest.
Change-Id: I744b624396fb8e8450529273cf4caae8476a5725
diff --git a/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/TaskManager.kt b/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/TaskManager.kt
index 1c18aa9..3dbf980 100644
--- a/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/TaskManager.kt
+++ b/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/TaskManager.kt
@@ -49,6 +49,7 @@
 import com.android.build.gradle.internal.component.TestComponentCreationConfig
 import com.android.build.gradle.internal.component.TestCreationConfig
 import com.android.build.gradle.internal.component.VariantCreationConfig
+import com.android.build.gradle.internal.coverage.CoverageAgentConfigurations
 import com.android.build.gradle.internal.coverage.JacocoConfigurations
 import com.android.build.gradle.internal.coverage.JacocoPropertiesTask
 import com.android.build.gradle.internal.coverage.JacocoReportTask
@@ -1581,16 +1582,28 @@
     }
   }
 
+  /**
+   * Maven coordinate of the on-device coverage agent for this project, computed (and its version validated) once: the agent version is the
+   * AGP version unless pinned with [StringOption.COVERAGE_AGENT_VERSION_OVERRIDE].
+   */
+  private val onTheFlyCoverageAgentDependency: String by lazy {
+    CoverageAgentConfigurations.getAgentRuntimeDependency(
+      globalConfig.services.projectOptions.get(StringOption.COVERAGE_AGENT_VERSION_OVERRIDE),
+      globalConfig.services.issueReporter,
+    )
+  }
+
   private fun handleOnTheFlyCoverageDependencies(creationConfig: ComponentCreationConfig) {
     val onTheFlyEnabled = globalConfig.services.projectOptions.get(BooleanOption.ENABLE_ON_THE_FLY_CODE_COVERAGE)
     if (onTheFlyEnabled && creationConfig.componentType.isForTesting && creationConfig.componentType.isApk) {
-      val coverageAgentDependency = com.android.build.gradle.internal.coverage.CoverageAgentConfigurations.getAgentRuntimeDependency()
+      val coverageAgentDependency = onTheFlyCoverageAgentDependency
       project.dependencies.add(creationConfig.variantDependencies.runtimeClasspath.name, coverageAgentDependency)
 
-      // The agent is released in lockstep with AGP, so force its version on the test component's
-      // classpaths: a different version declared by the build author (or pulled in transitively)
-      // must not win conflict resolution on the runtime classpath, nor leave the compile classpath
-      // on another version. Forcing is a no-op where the module is not present.
+      // The agent version is the AGP version (or the release pinned via the
+      // StringOption.COVERAGE_AGENT_VERSION_OVERRIDE gradle property), so force it on the test
+      // component's classpaths: a different version declared by the build author (or pulled in
+      // transitively) must not win conflict resolution on the runtime classpath, nor leave the
+      // compile classpath on another version. Forcing is a no-op where the module is not present.
       listOf(creationConfig.variantDependencies.compileClasspath, creationConfig.variantDependencies.runtimeClasspath).forEach {
         it.resolutionStrategy { r: ResolutionStrategy -> r.force(coverageAgentDependency) }
       }
diff --git a/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurations.kt b/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurations.kt
index 7b1f55f..d90a248 100644
--- a/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurations.kt
+++ b/build-system/gradle-core/src/main/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurations.kt
@@ -17,6 +17,9 @@
 package com.android.build.gradle.internal.coverage
 
 import com.android.Version
+import com.android.build.gradle.options.StringOption
+import com.android.builder.errors.IssueReporter
+import com.android.ide.common.repository.AgpVersion
 
 /** Configurations for the on-device coverage agent. */
 object CoverageAgentConfigurations {
@@ -26,13 +29,57 @@
   /**
    * Returns the Maven coordinate of the on-device coverage agent.
    *
-   * The agent is built and released in lockstep with the Android Gradle plugin (same version, same gmaven release), so it always resolves
-   * to the exact AGP version: the coordinate is added to the androidTest runtime classpath and forced on both the compile and runtime
-   * classpaths of the test component (see `TaskManager.handleOnTheFlyCoverageDependencies`), so a different version declared by the build
-   * author cannot win conflict resolution. Fixes to the agent ship with the next AGP release.
+   * The agent is built and released in lockstep with the Android Gradle plugin (same version, same gmaven release), so by default it
+   * resolves to the exact AGP version. A project can pin a different release of the agent with the
+   * [StringOption.COVERAGE_AGENT_VERSION_OVERRIDE] gradle property ([versionOverride]), e.g. to pick up an agent fix without changing its
+   * AGP version.
+   *
+   * Either way, the coordinate is added to the androidTest runtime classpath and forced on both the compile and runtime classpaths of the
+   * test component (see `TaskManager.handleOnTheFlyCoverageDependencies`), so a different version declared by the build author cannot win
+   * conflict resolution.
    */
   @JvmStatic
-  fun getAgentRuntimeDependency(): String {
-    return "$AGENT_DEPENDENCY_GROUP:$AGENT_DEPENDENCY_NAME:${Version.ANDROID_GRADLE_PLUGIN_VERSION}"
+  fun getAgentRuntimeDependency(versionOverride: String?, issueReporter: IssueReporter?): String {
+    return "$AGENT_DEPENDENCY_GROUP:$AGENT_DEPENDENCY_NAME:${getAgentVersion(versionOverride, issueReporter)}"
+  }
+
+  /**
+   * Returns the version of the on-device coverage agent to use: [versionOverride] if it is set and valid, otherwise [agpVersion]. Problems
+   * with the override are reported to [issueReporter] as a side effect.
+   */
+  @JvmStatic
+  @JvmOverloads
+  fun getAgentVersion(
+    versionOverride: String?,
+    issueReporter: IssueReporter?,
+    agpVersion: String = Version.ANDROID_GRADLE_PLUGIN_VERSION,
+  ): String {
+    val override = versionOverride?.trim()
+    if (override.isNullOrEmpty()) {
+      return agpVersion
+    }
+    val parsedOverride = AgpVersion.tryParse(override)
+    if (parsedOverride == null) {
+      issueReporter?.reportError(
+        IssueReporter.Type.GENERIC,
+        """
+        Could not parse on-device coverage agent version override '$override'.
+        Recommendation: Remove or fix the gradle property ${StringOption.COVERAGE_AGENT_VERSION_OVERRIDE.propertyName}; without it the agent version defaults to the Android Gradle plugin version ($agpVersion).
+        """
+          .trimIndent(),
+      )
+      return agpVersion
+    }
+    if (parsedOverride < AgpVersion.parse(agpVersion)) {
+      issueReporter?.reportWarning(
+        IssueReporter.Type.GENERIC,
+        """
+        The on-device coverage agent version override '$override' is older than the Android Gradle plugin version ($agpVersion).
+        Remove the gradle property ${StringOption.COVERAGE_AGENT_VERSION_OVERRIDE.propertyName} to use the agent released with the plugin.
+        """
+          .trimIndent(),
+      )
+    }
+    return override
   }
 }
diff --git a/build-system/gradle-core/src/main/java/com/android/build/gradle/options/StringOption.kt b/build-system/gradle-core/src/main/java/com/android/build/gradle/options/StringOption.kt
index ecdf438..394f28b 100644
--- a/build-system/gradle-core/src/main/java/com/android/build/gradle/options/StringOption.kt
+++ b/build-system/gradle-core/src/main/java/com/android/build/gradle/options/StringOption.kt
@@ -145,6 +145,15 @@
   // An override for the Jacoco version to be set per module
   JACOCO_TOOL_VERSION("com.android.test.jacocoVersion", ApiStage.Experimental),
 
+  /**
+   * Overrides the version of the on-device coverage agent (com.android.tools.test:coverage-agent) used when
+   * [BooleanOption.ENABLE_ON_THE_FLY_CODE_COVERAGE] is enabled.
+   *
+   * The agent is released in lockstep with the Android Gradle plugin and defaults to the plugin's version. Set this to pick up the agent
+   * from a different release without changing the plugin version.
+   */
+  COVERAGE_AGENT_VERSION_OVERRIDE("android.experimental.androidTest.coverage.onDeviceAgentVersion", ApiStage.Experimental),
+
   /** If a value is set, all version-gated [Option] with a minimum AGP version less than or equal to this value will be enabled. */
   SIMULATE_AGP_VERSION_BEHAVIOR("android.simulateAgpVersionBehavior", ApiStage.Experimental),
 
diff --git a/build-system/gradle-core/src/test/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurationsTest.kt b/build-system/gradle-core/src/test/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurationsTest.kt
index 3b91a2d..4280b3c 100644
--- a/build-system/gradle-core/src/test/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurationsTest.kt
+++ b/build-system/gradle-core/src/test/java/com/android/build/gradle/internal/coverage/CoverageAgentConfigurationsTest.kt
@@ -17,19 +17,94 @@
 package com.android.build.gradle.internal.coverage
 
 import com.android.Version
+import com.android.build.gradle.internal.fixtures.FakeSyncIssueReporter
+import com.android.build.gradle.options.StringOption
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
 
 class CoverageAgentConfigurationsTest {
 
+  private val agpVersion = "9.5.0-alpha09"
+  private val propertyName = StringOption.COVERAGE_AGENT_VERSION_OVERRIDE.propertyName
+  private val issueReporter = FakeSyncIssueReporter(throwOnError = false)
+
   @Test
-  fun agentCoordinateIsPinnedToTheAgpVersion() {
-    val coordinate = CoverageAgentConfigurations.getAgentRuntimeDependency()
+  fun agentCoordinateIsPinnedToTheAgpVersionByDefault() {
+    val coordinate = CoverageAgentConfigurations.getAgentRuntimeDependency(versionOverride = null, issueReporter = issueReporter)
 
     assertThat(coordinate).isEqualTo("com.android.tools.test:coverage-agent:${Version.ANDROID_GRADLE_PLUGIN_VERSION}")
     // The agent is released in lockstep with AGP: it must never carry an independent version
     // (it used to be 0.1.0-alpha01 before it was coupled to the plugin).
     assertThat(coordinate.substringAfterLast(':')).isEqualTo(Version.ANDROID_GRADLE_PLUGIN_VERSION)
     assertThat(coordinate).doesNotContain("0.1.0")
+    assertThat(issueReporter.messages).isEmpty()
+  }
+
+  @Test
+  fun agentCoordinateUsesTheOverriddenVersion() {
+    val coordinate = CoverageAgentConfigurations.getAgentRuntimeDependency(versionOverride = "99.0.0", issueReporter = issueReporter)
+
+    assertThat(coordinate).isEqualTo("com.android.tools.test:coverage-agent:99.0.0")
+    assertThat(issueReporter.messages).isEmpty()
+  }
+
+  @Test
+  fun blankOverrideUsesTheAgpVersion() {
+    for (override in listOf(null, "", "   ")) {
+      assertThat(CoverageAgentConfigurations.getAgentVersion(override, issueReporter, agpVersion)).isEqualTo(agpVersion)
+    }
+    assertThat(issueReporter.messages).isEmpty()
+  }
+
+  @Test
+  fun newerOverrideIsUsed() {
+    assertThat(CoverageAgentConfigurations.getAgentVersion("9.5.0-alpha12", issueReporter, agpVersion)).isEqualTo("9.5.0-alpha12")
+    assertThat(CoverageAgentConfigurations.getAgentVersion("9.6.0", issueReporter, agpVersion)).isEqualTo("9.6.0")
+    assertThat(issueReporter.messages).isEmpty()
+  }
+
+  @Test
+  fun sameVersionOverrideIsUsedSilently() {
+    assertThat(CoverageAgentConfigurations.getAgentVersion(agpVersion, issueReporter, agpVersion)).isEqualTo(agpVersion)
+    assertThat(issueReporter.messages).isEmpty()
+  }
+
+  @Test
+  fun overrideIsTrimmed() {
+    assertThat(CoverageAgentConfigurations.getAgentVersion("  9.5.0-alpha12 ", issueReporter, agpVersion)).isEqualTo("9.5.0-alpha12")
+    assertThat(issueReporter.messages).isEmpty()
+  }
+
+  @Test
+  fun olderOverrideIsUsedWithAWarning() {
+    assertThat(CoverageAgentConfigurations.getAgentVersion("9.4.2", issueReporter, agpVersion)).isEqualTo("9.4.2")
+
+    assertThat(issueReporter.errors).isEmpty()
+    val warning = issueReporter.warnings.single()
+    assertThat(warning).contains("'9.4.2' is older than the Android Gradle plugin version (9.5.0-alpha09)")
+    assertThat(warning).contains(propertyName)
+  }
+
+  @Test
+  fun olderPreviewOfTheSameSeriesIsUsedWithAWarning() {
+    assertThat(CoverageAgentConfigurations.getAgentVersion("9.5.0-alpha08", issueReporter, agpVersion)).isEqualTo("9.5.0-alpha08")
+
+    assertThat(issueReporter.errors).isEmpty()
+    assertThat(issueReporter.warnings).hasSize(1)
+  }
+
+  @Test
+  fun unparsableOverrideIsReportedAndFallsBackToTheAgpVersion() {
+    for (override in listOf("9.5", "latest", "+", "9.5.0-SNAPSHOT")) {
+      val fakeIssueReporter = FakeSyncIssueReporter(throwOnError = false)
+
+      assertThat(CoverageAgentConfigurations.getAgentVersion(override, fakeIssueReporter, agpVersion)).isEqualTo(agpVersion)
+
+      assertThat(fakeIssueReporter.warnings).isEmpty()
+      val error = fakeIssueReporter.errors.single()
+      assertThat(error).contains("Could not parse on-device coverage agent version override '$override'")
+      assertThat(error).contains("Remove or fix the gradle property $propertyName")
+      assertThat(error).contains(agpVersion)
+    }
   }
 }
diff --git a/build-system/integration-test/application/src/test/java/com/android/build/gradle/integration/application/JacocoTest.java b/build-system/integration-test/application/src/test/java/com/android/build/gradle/integration/application/JacocoTest.java
index 9129f4a..0965557 100644
--- a/build-system/integration-test/application/src/test/java/com/android/build/gradle/integration/application/JacocoTest.java
+++ b/build-system/integration-test/application/src/test/java/com/android/build/gradle/integration/application/JacocoTest.java
@@ -28,6 +28,8 @@
 import com.android.build.gradle.integration.common.utils.TestFileUtils;
 import com.android.build.gradle.internal.coverage.JacocoConfigurations;
 import com.android.build.gradle.options.BooleanOption;
+import com.android.build.gradle.options.StringOption;
+import com.android.testutils.MavenRepoGenerator;
 import com.android.testutils.TestInputsGenerator;
 import com.android.testutils.apk.Apk;
 import com.android.testutils.apk.Dex;
@@ -70,6 +72,7 @@
         this.project = GradleTestProject.builder()
                 .fromTestApp(TEST_APP)
                 .addGradleProperty(BooleanOption.ANDROID_BUILTIN_TEST_PLATFORM, runWithBuiltInPlatform)
+                .withAdditionalMavenRepo(fakeAgentRepo())
                 .create();
     }
 
@@ -84,6 +87,35 @@
     private static final GradleProject TEST_APP =
             HelloWorldApp.forPlugin("com.android.application");
 
+    private static final String AGENT_COORDINATE_PREFIX = "com.android.tools.test:coverage-agent:";
+    private static final String AGENT_ATTACHER_CLASS =
+            "Lcom/android/tools/coverage/CoverageAgentAttacher;";
+
+    // A fake release of the on-device coverage agent, newer than any real AGP version, that exists
+    // only in the test project's additional Maven repository. It contains a marker class instead
+    // of the real attacher, so the content of an androidTest APK built against it tells which
+    // release was used.
+    private static final String NEWER_FAKE_AGENT_VERSION = "99.0.0";
+    private static final String FAKE_AGENT_MARKER_CLASS_NAME =
+            "com/android/tools/coverage/FakeCoverageAgentMarker";
+    private static final String FAKE_AGENT_MARKER_CLASS_TYPE =
+            "L" + FAKE_AGENT_MARKER_CLASS_NAME + ";";
+
+    private static MavenRepoGenerator fakeAgentRepo() {
+        try {
+            byte[] fakeAgentJar =
+                    TestInputsGenerator.jarWithEmptyClasses(
+                            Collections.singleton(FAKE_AGENT_MARKER_CLASS_NAME));
+            return new MavenRepoGenerator(
+                    Collections.singletonList(
+                            new MavenRepoGenerator.Library(
+                                    AGENT_COORDINATE_PREFIX + NEWER_FAKE_AGENT_VERSION,
+                                    fakeAgentJar)));
+        } catch (Exception e) {
+            throw new IllegalStateException("Could not generate the fake coverage agent jar", e);
+        }
+    }
+
     private File getCoverageDir() {
         return runWithBuiltInPlatform
                 ? new File(project.getProjectDir(), "build/intermediates/test_suite_code_coverage/debug/connectedDebugAndroidTest")
@@ -527,6 +559,87 @@
     }
 
     @Test
+    public void onTheFlyAgentVersionCanBeOverriddenWithAGradleProperty() throws Exception {
+        String staleCoordinate = AGENT_COORDINATE_PREFIX + "0.1.0-alpha01";
+        TestFileUtils.appendToFile(
+                project.getBuildFile(),
+                "\ndependencies {\n"
+                        + "    androidTestImplementation '"
+                        + staleCoordinate
+                        + "'\n"
+                        + "}\n");
+
+        // A project can pin the agent to another release than the one bundled with its AGP
+        // version, e.g. to pick up an agent fix without upgrading the plugin. The override, not
+        // the AGP version, is then what both androidTest classpaths resolve, and the force follows
+        // it: the stale version declared by the build author is substituted with the override...
+        for (String configuration :
+                Arrays.asList(
+                        "debugAndroidTestCompileClasspath", "debugAndroidTestRuntimeClasspath")) {
+            String agentReport = agentDependencyReport(configuration, NEWER_FAKE_AGENT_VERSION);
+            assertThat(agentReport)
+                    .named(configuration)
+                    .contains(staleCoordinate + " -> " + NEWER_FAKE_AGENT_VERSION);
+            assertAgentResolvedTo(agentReport, configuration, NEWER_FAKE_AGENT_VERSION);
+        }
+
+        // ... and what gets packaged into the androidTest APK: the fake release only has the
+        // marker class, the real one only has the attacher.
+        project.executor()
+                .with(BooleanOption.ENABLE_ON_THE_FLY_CODE_COVERAGE, true)
+                .with(StringOption.COVERAGE_AGENT_VERSION_OVERRIDE, NEWER_FAKE_AGENT_VERSION)
+                .run("assembleDebugAndroidTest");
+        try (Apk androidTestApk = project.getApk(GradleTestProject.ApkType.ANDROIDTEST_DEBUG)) {
+            assertThat(androidTestApk).containsClass(FAKE_AGENT_MARKER_CLASS_TYPE);
+            assertThat(androidTestApk).doesNotContainClass(AGENT_ATTACHER_CLASS);
+        }
+    }
+
+    /**
+     * Runs the dependency report of the given configuration with on-the-fly coverage enabled and
+     * the agent version overridden, and returns the lines of the report that mention the agent.
+     */
+    private String agentDependencyReport(String configuration, String agentVersionOverride)
+            throws Exception {
+        GradleBuildResult result =
+                project.executor()
+                        .with(BooleanOption.ENABLE_ON_THE_FLY_CODE_COVERAGE, true)
+                        .with(StringOption.COVERAGE_AGENT_VERSION_OVERRIDE, agentVersionOverride)
+                        .run(":dependencies", "--configuration", configuration);
+        StringBuilder agentLines = new StringBuilder();
+        try (Scanner scanner = result.getStdout()) {
+            while (scanner.hasNextLine()) {
+                String line = scanner.nextLine();
+                if (line.contains(AGENT_COORDINATE_PREFIX)) {
+                    agentLines.append(line).append('\n');
+                }
+            }
+        }
+        assertThat(agentLines.toString()).named("coverage agent on " + configuration).isNotEmpty();
+        return agentLines.toString();
+    }
+
+    /**
+     * Asserts that every agent line of a dependency report resolves to the given version, whether
+     * the agent is requested directly ("group:name:version") or substituted ("group:name:requested
+     * -> resolved").
+     */
+    private static void assertAgentResolvedTo(
+            String agentReport, String configuration, String version) {
+        for (String line : agentReport.split("\n")) {
+            String resolved =
+                    line.contains(" -> ")
+                            ? line.substring(line.indexOf(" -> ") + 4)
+                            : line.substring(
+                                    line.indexOf(AGENT_COORDINATE_PREFIX)
+                                            + AGENT_COORDINATE_PREFIX.length());
+            assertThat(resolved.trim().split(" ")[0])
+                    .named(configuration + ": " + line)
+                    .isEqualTo(version);
+        }
+    }
+
+    @Test
     public void checkJacocoAntConfiguration() throws IOException, InterruptedException {
         // jacoco ant configuration exists by default due to the usage from
         // TestResultsCollectionTask for the aggregated test reporting feature