Clean up JVM test setup where classes from aars are used as input

Introduce androidx.enableAarAsJarForJvmTest() that adds a new
configuration testAarAsJar that allows annotation project JVM tests
to add AAR classes to JVM test classpath.

Test: ./gradlew test
Change-Id: I4d24c2946da99c5762947c26f299860baf38dc26
diff --git a/appsearch/appsearch/build.gradle b/appsearch/appsearch/build.gradle
index 02e32f9..c9e4a8b 100644
--- a/appsearch/appsearch/build.gradle
+++ b/appsearch/appsearch/build.gradle
@@ -50,17 +50,6 @@
     androidTestImplementation(libs.protobufLite)
 }
 
-// Create a jar for use by appsearch-compiler:test
-android.libraryVariants.all { variant ->
-    def name = variant.name
-    def suffix = name.capitalize()
-    project.tasks.create(name: "jar${suffix}", type: Jar) {
-        dependsOn variant.javaCompileProvider.get()
-        from variant.javaCompileProvider.get().destinationDirectory
-        destinationDirectory.set(new File(project.buildDir, "libJar"))
-    }
-}
-
 androidx {
     name = 'AndroidX AppSearch'
     publish = Publish.SNAPSHOT_AND_RELEASE
diff --git a/appsearch/compiler/build.gradle b/appsearch/compiler/build.gradle
index 486ae37..3aee826 100644
--- a/appsearch/compiler/build.gradle
+++ b/appsearch/compiler/build.gradle
@@ -21,6 +21,8 @@
     id('java-library')
 }
 
+androidx.enableAarAsJarForJvmTest()
+
 dependencies {
     api('androidx.annotation:annotation:1.1.0')
     api(libs.jsr250)
@@ -30,19 +32,10 @@
     implementation(libs.javapoet)
 
     // For testing, add in the compiled classes from appsearch to get access to annotations.
-    testImplementation fileTree(
-            dir: provider {
-                // Wrapping in a provider as a workaround as we access buildDir before this project is configured
-                // Replace with AGP API once it is added b/228109260
-                "${new File(project(":appsearch:appsearch").buildDir, "libJar")}"
-            },
-            include : "*.jar"
-    )
+    testAarAsJar(project(":appsearch:appsearch"))
     testImplementation(libs.googleCompileTesting)
 }
 
-tasks.findByName('compileJava').dependsOn(":appsearch:appsearch:jarRelease")
-
 androidx {
     name = 'AndroidX AppSearch Compiler'
     type = LibraryType.ANNOTATION_PROCESSOR
diff --git a/buildSrc/private/src/main/kotlin/androidx/build/AndroidXExtension.kt b/buildSrc/private/src/main/kotlin/androidx/build/AndroidXExtension.kt
index c12c696..ac3da1b 100644
--- a/buildSrc/private/src/main/kotlin/androidx/build/AndroidXExtension.kt
+++ b/buildSrc/private/src/main/kotlin/androidx/build/AndroidXExtension.kt
@@ -17,6 +17,7 @@
 package androidx.build
 
 import androidx.build.checkapi.shouldConfigureApiTasks
+import androidx.build.jvmtest.configureAarAsJarForJvmTest
 import com.android.build.gradle.internal.crash.afterEvaluate
 import groovy.lang.Closure
 import org.gradle.api.GradleException
@@ -248,6 +249,10 @@
         return licenses
     }
 
+    fun enableAarAsJarForJvmTest() {
+        configureAarAsJarForJvmTest(project)
+    }
+
     companion object {
         const val DEFAULT_UNSPECIFIED_VERSION = "unspecified"
     }
diff --git a/buildSrc/private/src/main/kotlin/androidx/build/jvmtest/AarDependencyForJvmTest.kt b/buildSrc/private/src/main/kotlin/androidx/build/jvmtest/AarDependencyForJvmTest.kt
new file mode 100644
index 0000000..d9f45ec
--- /dev/null
+++ b/buildSrc/private/src/main/kotlin/androidx/build/jvmtest/AarDependencyForJvmTest.kt
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.build.jvmtest
+
+import androidx.build.transform.ExtractClassesJarTransform
+import androidx.build.transform.IdentityTransform
+import com.android.build.api.attributes.BuildTypeAttr
+import org.gradle.api.Project
+import org.gradle.api.attributes.Attribute
+import org.gradle.api.attributes.Usage
+
+/**
+ * Creates `testAarAsJar` configuration that can be used for JVM tests that need to Android library
+ * classes on the classpath.
+ */
+fun configureAarAsJarForJvmTest(project: Project) {
+    val testAarsAsJars = project.configurations.create("testAarAsJar") {
+        it.isTransitive = false
+        it.isCanBeConsumed = false
+        it.isCanBeResolved = true
+        it.attributes.attribute(
+            BuildTypeAttr.ATTRIBUTE,
+            project.objects.named(BuildTypeAttr::class.java, "release")
+        )
+        it.attributes.attribute(
+            Usage.USAGE_ATTRIBUTE,
+            project.objects.named(Usage::class.java, Usage.JAVA_API)
+        )
+    }
+    val artifactType = Attribute.of("artifactType", String::class.java)
+    project.dependencies.registerTransform(IdentityTransform::class.java) { spec ->
+        spec.from.attribute(artifactType, "jar")
+        spec.to.attribute(artifactType, "aarAsJar")
+    }
+
+    project.dependencies.registerTransform(ExtractClassesJarTransform::class.java) { spec ->
+        spec.from.attribute(artifactType, "aar")
+        spec.to.attribute(artifactType, "aarAsJar")
+    }
+
+    val aarAsJar = testAarsAsJars.incoming.artifactView { viewConfiguration ->
+        viewConfiguration.attributes.attribute(artifactType, "aarAsJar")
+    }.files
+    project.configurations.getByName("testImplementation").dependencies.add(
+        project.dependencies.create(aarAsJar)
+    )
+}
\ No newline at end of file
diff --git a/buildSrc/private/src/main/kotlin/androidx/build/transform/ExtractClassesJarTransform.kt b/buildSrc/private/src/main/kotlin/androidx/build/transform/ExtractClassesJarTransform.kt
new file mode 100644
index 0000000..cfeff81
--- /dev/null
+++ b/buildSrc/private/src/main/kotlin/androidx/build/transform/ExtractClassesJarTransform.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.build.transform
+
+import com.google.common.io.Files
+import java.util.zip.ZipInputStream
+import org.gradle.api.artifacts.transform.InputArtifact
+import org.gradle.api.artifacts.transform.TransformAction
+import org.gradle.api.artifacts.transform.TransformOutputs
+import org.gradle.api.artifacts.transform.TransformParameters
+import org.gradle.api.file.FileSystemLocation
+import org.gradle.api.provider.Provider
+import org.gradle.api.tasks.PathSensitive
+import org.gradle.api.tasks.PathSensitivity
+import org.gradle.work.DisableCachingByDefault
+
+@DisableCachingByDefault
+abstract class ExtractClassesJarTransform : TransformAction<TransformParameters.None> {
+    @get:PathSensitive(PathSensitivity.NAME_ONLY)
+    @get:InputArtifact
+    abstract val primaryInput: Provider<FileSystemLocation>
+
+    override fun transform(outputs: TransformOutputs) {
+        val inputFile = primaryInput.get().asFile
+        val outputFile = outputs.file("${inputFile.nameWithoutExtension}.jar")
+        ZipInputStream(inputFile.inputStream().buffered()).use { zipInputStream ->
+            while (true) {
+                val entry = zipInputStream.nextEntry ?: break
+                if (entry.name != "classes.jar") continue
+                Files.asByteSink(outputFile).writeFrom(zipInputStream)
+                break
+            }
+        }
+    }
+}
diff --git a/buildSrc/private/src/main/kotlin/androidx/build/transform/IdentityTransform.kt b/buildSrc/private/src/main/kotlin/androidx/build/transform/IdentityTransform.kt
new file mode 100644
index 0000000..d662963
--- /dev/null
+++ b/buildSrc/private/src/main/kotlin/androidx/build/transform/IdentityTransform.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.build.transform
+
+import org.gradle.api.artifacts.transform.InputArtifact
+import org.gradle.api.artifacts.transform.TransformAction
+import org.gradle.api.artifacts.transform.TransformOutputs
+import org.gradle.api.artifacts.transform.TransformParameters
+import org.gradle.api.file.FileSystemLocation
+import org.gradle.api.provider.Provider
+import org.gradle.api.tasks.PathSensitive
+import org.gradle.api.tasks.PathSensitivity
+import org.gradle.work.DisableCachingByDefault
+
+@DisableCachingByDefault
+abstract class IdentityTransform : TransformAction<TransformParameters.None> {
+    @get:PathSensitive(PathSensitivity.ABSOLUTE)
+    @get:InputArtifact
+    abstract val inputArtifact: Provider<FileSystemLocation>
+
+    override fun transform(transformOutputs: TransformOutputs) {
+        val input = inputArtifact.get().asFile
+        when {
+            input.isDirectory -> transformOutputs.dir(input)
+            input.isFile -> transformOutputs.file(input)
+            else -> throw IllegalArgumentException(
+                "File/directory does not exist: ${input.absolutePath}")
+        }
+    }
+}
\ No newline at end of file
diff --git a/hilt/hilt-compiler/build.gradle b/hilt/hilt-compiler/build.gradle
index 903b3be..fd43c06 100644
--- a/hilt/hilt-compiler/build.gradle
+++ b/hilt/hilt-compiler/build.gradle
@@ -24,6 +24,8 @@
     id("kotlin-kapt")
 }
 
+androidx.enableAarAsJarForJvmTest()
+
 dependencies {
     implementation(libs.kotlinStdlib)
     compileOnly(libs.autoServiceAnnotations)
@@ -38,22 +40,12 @@
     testImplementation(libs.truth)
     testImplementation(libs.googleCompileTesting)
     testImplementation(libs.hiltCore)
-    testImplementation(fileTree(
-            dir: provider {
-                // Replace with AGP API once it is added b/228109260
-                // Wrapping in a provider as a workaround as we access buildDir before this project is configured
-                "${new File(project(":hilt:hilt-work").buildDir, "libJar")}"
-            },
-            include : "*.jar"))
+    testAarAsJar(project(":hilt:hilt-work"))
     testImplementation(fileTree(
             dir: "${SdkHelperKt.getSdkPath(project)}/platforms/$SupportConfig.COMPILE_SDK_VERSION/",
             include : "android.jar"))
 }
 
-tasks.named("compileKotlin").configure {
-    dependsOn(":hilt:hilt-work:jarRelease")
-}
-
 androidx {
     name = "AndroidX Hilt Extension Compiler"
     type = LibraryType.ANNOTATION_PROCESSOR
diff --git a/hilt/hilt-work/build.gradle b/hilt/hilt-work/build.gradle
index ec2e644..da8f26b 100644
--- a/hilt/hilt-work/build.gradle
+++ b/hilt/hilt-work/build.gradle
@@ -36,18 +36,6 @@
     annotationProcessor(libs.hiltCompiler)
 }
 
-android.libraryVariants.all { variant ->
-    def name = variant.name
-    def suffix = name.capitalize()
-
-    // Create jar<variant> task for testImplementation in hilt-compiler.
-    project.tasks.register("jar${suffix}", Jar).configure {
-        dependsOn(variant.javaCompileProvider)
-        from(variant.javaCompileProvider.map { task -> task.destinationDir})
-        destinationDirectory.fileValue(new File(project.buildDir, "libJar"))
-    }
-}
-
 androidx {
     name = "Android Lifecycle WorkManager Hilt Extension"
     publish = Publish.SNAPSHOT_AND_RELEASE
diff --git a/lifecycle/lifecycle-viewmodel-savedstate/build.gradle b/lifecycle/lifecycle-viewmodel-savedstate/build.gradle
index 8a3544c..2f5146a 100644
--- a/lifecycle/lifecycle-viewmodel-savedstate/build.gradle
+++ b/lifecycle/lifecycle-viewmodel-savedstate/build.gradle
@@ -14,9 +14,7 @@
  * limitations under the License.
  */
 
-
 import androidx.build.Publish
-import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 
 plugins {
     id("AndroidXPlugin")
@@ -61,19 +59,6 @@
     }
 }
 
-//used by testImplementation safe-args-generator
-android.libraryVariants.all { variant ->
-    def name = variant.name
-    def suffix = name.capitalize()
-    project.tasks.register("jar${suffix}", Copy).configure {
-        dependsOn ("assemble$suffix")
-        from(zipTree("${project.buildDir}/outputs/aar/lifecycle-viewmodel-savedstate-${name}.aar")) {
-            include "classes.jar"
-        }
-        destinationDir new File(project.buildDir, "libJar")
-    }
-}
-
 androidx {
     name = "Android Lifecycle ViewModel with SavedState"
     publish = Publish.SNAPSHOT_AND_RELEASE
diff --git a/navigation/navigation-common/build.gradle b/navigation/navigation-common/build.gradle
index 8bb9c03..1428843 100644
--- a/navigation/navigation-common/build.gradle
+++ b/navigation/navigation-common/build.gradle
@@ -15,7 +15,6 @@
  */
 
 import androidx.build.Publish
-import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 
 plugins {
     id("AndroidXPlugin")
@@ -60,19 +59,6 @@
     lintPublish(project(':navigation:navigation-common-lint'))
 }
 
-//used by testImplementation safe-args-generator
-android.libraryVariants.all { variant ->
-    def name = variant.name
-    def suffix = name.capitalize()
-    project.tasks.register("jar${suffix}", Copy).configure {
-        dependsOn("assemble$suffix")
-        from(zipTree("${project.buildDir}/outputs/aar/navigation-common-${name}.aar")) {
-            include("classes.jar")
-        }
-        destinationDir(new File(project.buildDir, "libJar"))
-    }
-}
-
 androidx {
     name = "Android Navigation Common"
     publish = Publish.SNAPSHOT_AND_RELEASE
diff --git a/navigation/navigation-safe-args-generator/build.gradle b/navigation/navigation-safe-args-generator/build.gradle
index 0a72090..1b6930c 100644
--- a/navigation/navigation-safe-args-generator/build.gradle
+++ b/navigation/navigation-safe-args-generator/build.gradle
@@ -23,6 +23,8 @@
     id("kotlin")
 }
 
+androidx.enableAarAsJarForJvmTest()
+
 dependencies {
     implementation(libs.xpp3)
     implementation(libs.xmlpull)
@@ -40,22 +42,8 @@
             dir: "${SdkHelperKt.getSdkPath(project)}/platforms/$SupportConfig.COMPILE_SDK_VERSION/",
             include : "android.jar"
     ))
-    testImplementation(fileTree(
-            dir: provider {
-                // Wrapping in a provider as a workaround as we access buildDir before this project is configured
-                // Replace with AGP API once it is added b/228109260
-                "${new File(project(":navigation:navigation-common").buildDir, "libJar")}"
-            },
-            include : "*.jar"
-    ))
-    testImplementation(fileTree(
-            dir: provider {
-                // Wrapping in a provider as a workaround as we access buildDir before this project is configured
-                // Replace with AGP API once it is added b/228109260
-                "${new File(project(":lifecycle:lifecycle-viewmodel-savedstate").buildDir, "libJar")}"
-            },
-            include : "*.jar"
-    ))
+    testAarAsJar(project(":navigation:navigation-common"))
+    testAarAsJar(project(":lifecycle:lifecycle-viewmodel-savedstate"))
 }
 
 tasks.findByName("test").doFirst {
@@ -66,9 +54,6 @@
     it.classpath = files(classpath.minus(androidJar).plus(androidJar))
 }
 
-tasks.findByName("compileKotlin").dependsOn(":navigation:navigation-common:jarRelease")
-tasks.findByName("compileKotlin").dependsOn(":lifecycle:lifecycle-viewmodel-savedstate:jarRelease")
-
 androidx {
     name = "Android Navigation TypeSafe Arguments Generator"
     type = LibraryType.OTHER_CODE_PROCESSOR
diff --git a/room/room-compiler/build.gradle b/room/room-compiler/build.gradle
index 3edc9a10..6ead8ab 100644
--- a/room/room-compiler/build.gradle
+++ b/room/room-compiler/build.gradle
@@ -84,6 +84,8 @@
     }
 }
 
+androidx.enableAarAsJarForJvmTest()
+
 dependencies {
     implementation(project(":room:room-common"))
     implementation(project(":room:room-migration"))
@@ -113,22 +115,8 @@
             dir: "${SdkHelperKt.getSdkPath(project)}/platforms/$SupportConfig.COMPILE_SDK_VERSION/",
             include : "android.jar"
     ))
-    testImplementation(fileTree(
-            dir: provider {
-                // Wrapping in a provider as we access buildDir before this project is configured
-                // Replace with AGP API once it is added b/228109260
-                "${new File(project(":room:room-runtime").buildDir, "intermediates/runtime_library_classes_jar/release/")}"
-            },
-            include : "*.jar"
-    ))
-    testImplementation(fileTree(
-            dir: provider {
-                // Wrapping in a provider as we access buildDir before this project is configured
-                // Replace with AGP API once it is added b/228109260
-                "${new File(project(":sqlite:sqlite").buildDir, "intermediates/compile_library_classes_jar/release/")}"
-            },
-            include : "*.jar"
-    ))
+    testAarAsJar(project(":room:room-runtime"))
+    testAarAsJar(project(":sqlite:sqlite"))
     testImplementation(project(":internal-testutils-common"))
 }
 
@@ -284,8 +272,6 @@
 
 tasks.findByName("compileKotlin").dependsOn(generateAntlrTask)
 tasks.findByName("sourceJar").dependsOn(generateAntlrTask)
-tasks.findByName("compileKotlin").dependsOn(":room:room-runtime:bundleLibRuntimeToJarRelease")
-tasks.findByName("compileKotlin").dependsOn(":sqlite:sqlite:jarRelease")
 
 tasks.withType(KotlinCompile).configureEach {
     kotlinOptions {
diff --git a/room/room-runtime/build.gradle b/room/room-runtime/build.gradle
index 56968c37..8eb962e 100644
--- a/room/room-runtime/build.gradle
+++ b/room/room-runtime/build.gradle
@@ -69,18 +69,6 @@
 
 }
 
-android.libraryVariants.all { variant ->
-    def name = variant.name
-    def suffix = name.capitalize()
-
-    // Create jar<variant> task for testImplementation in room-compiler.
-    project.tasks.create(name: "jar${suffix}", type: Jar){
-        dependsOn(variant.javaCompileProvider.get())
-        from(variant.javaCompileProvider.get().destinationDir)
-        destinationDirectory.fileValue(new File(project.buildDir, "libJar"))
-    }
-}
-
 androidx {
     name = "Android Room-Runtime"
     publish = Publish.SNAPSHOT_AND_RELEASE
diff --git a/sqlite/sqlite/build.gradle b/sqlite/sqlite/build.gradle
index d3f6c60..ca04999 100644
--- a/sqlite/sqlite/build.gradle
+++ b/sqlite/sqlite/build.gradle
@@ -33,17 +33,6 @@
     namespace "androidx.sqlite.db"
 }
 
-// Used by testImplementation in room-compiler
-android.libraryVariants.all { variant ->
-    def name = variant.name
-    def suffix = name.capitalize()
-    def jarTask = project.tasks.create(name: "jar${suffix}", type: Jar){
-        dependsOn(variant.javaCompileProvider.get())
-        from(variant.javaCompileProvider.get().destinationDir)
-        destinationDirectory.fileValue(new File(project.buildDir, "libJar"))
-    }
-}
-
 androidx {
     name = "Android DB"
     publish = Publish.SNAPSHOT_AND_RELEASE