diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index dadbd9c..abb3e28 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -153,7 +153,7 @@ jobs:
 
       - name: Run instrumentation tests
         id: screenshotsverify
-        continue-on-error: true
+        continue-on-error: false
         uses: reactivecircus/android-emulator-runner@1dcd0090116d15e7c562f8db72807de5e036a4ed # v2.34.0
         with:
           api-level: 31
@@ -162,7 +162,7 @@ jobs:
           force-avd-creation: false
           emulator-options: -no-snapshot-save -no-window -gpu swiftshader_indirect -noaudio -no-boot-anim -camera-back none
           disable-animations: true
-          script: ./gradlew connectedCheck --stacktrace
+          script: ./gradlew connectedCheck --stacktrace && tests/gradlew -p tests :connectedCheck --stacktrace
 
       - name: Prevent pushing new screenshots if this is a fork
         id: checkfork_screenshots
@@ -176,9 +176,9 @@ jobs:
         continue-on-error: true
         if: steps.screenshotsverify.outcome == 'failure' && github.event_name == 'pull_request'
         run: |
-          echo "Pulling $(ls -1q dropshots/build/reports/androidTests/dropshots/reference/*.png | wc -l) files..."
-          ls -1q dropshots/build/reports/androidTests/dropshots/reference/*.png
-          cp dropshots/build/reports/androidTests/dropshots/reference/*.png dropshots/src/androidTest/assets/
+          echo "Pulling $(ls -1q tests/build/test-results/dropshots/reference/*.png | wc -l) files..."
+          ls -1q tests/build/test-results/dropshots/reference/*.png
+          cp tests/build/test-results/dropshots/reference/*.png tests/integration-test/src/androidTest/assets/
 
       # Since commits from actions don't trigger new actions, we validate the new screenshots here
       # before we commit them to ensure there isn't flakiness in the tests.
@@ -234,7 +234,7 @@ jobs:
         uses: gradle/actions/setup-gradle@v4
 
       - name: Upload Snapshot
-        run: ./gradlew publish -Pdropshots.releaseMode=true --no-daemon --no-parallel --no-configuration-cache --stacktrace
+        run: ./gradlew publish --no-daemon --no-parallel --no-configuration-cache --stacktrace
         env:
           ORG_GRADLE_PROJECT_mavenCentralUsername: ${{ secrets.SONATYPE_USERNAME }}
           ORG_GRADLE_PROJECT_mavenCentralPassword: ${{ secrets.SONATYPE_PASSWORD }}
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3804583..5fff2c0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,14 +3,14 @@
 ## [Unreleased]
 [Unreleased]: https://github.com/dropbox/dropshots/compare/0.5.0...HEAD
 
-New:
-- Nothing yet!
-
-Changed:
-- Nothing yet!
-
-Fixed:
-- Nothing yet!
+* Support for newDsl in Android Gradle Plugin 9.0.0
+* Changes default behavior of `recordDebugAndroidTestScreenshots` task to record screenshots even
+  if tests fail. This behavior can be disabled by setting `dropshots.recordOnFailure = false`
+  on the Gradle plugin configuration.
+* Exposed `Dropshots.recordScreenshots` as a public API to allow detecting whether screenshots are
+  being recorded in the current test run.
+* Added new optional argument to `Dropshots` constructor. `rootScreenshotDirectory` allows users to
+  specify a custom root directory for screenshots on the device.
 
 ## [0.5.0] = 2025-04-08
 [0.5.0]: https://github.com/dropbox/dropshots/releases/tags/0.5.0
diff --git a/dropshots-gradle-plugin/api/dropshots-gradle-plugin.api b/dropshots-gradle-plugin/api/dropshots-gradle-plugin.api
index 08d34b9..433f880 100644
--- a/dropshots-gradle-plugin/api/dropshots-gradle-plugin.api
+++ b/dropshots-gradle-plugin/api/dropshots-gradle-plugin.api
@@ -8,6 +8,7 @@ public abstract class com/dropbox/dropshots/ClearScreenshotsTask : org/gradle/ap
 
 public abstract class com/dropbox/dropshots/DropshotsExtension {
 	public fun <init> (Lorg/gradle/api/model/ObjectFactory;)V
+	public final fun getRecordOnFailure ()Lorg/gradle/api/provider/Property;
 	public final fun getReferenceOutputDirectory ()Lorg/gradle/api/provider/Property;
 }
 
@@ -21,8 +22,11 @@ public abstract class com/dropbox/dropshots/PullScreenshotsTask : org/gradle/api
 	public fun <init> ()V
 	public abstract fun getAdbExecutable ()Lorg/gradle/api/provider/Property;
 	protected abstract fun getExecOperations ()Lorg/gradle/process/ExecOperations;
+	protected abstract fun getFileOperations ()Lorg/gradle/api/file/FileSystemOperations;
 	public abstract fun getOutputDirectory ()Lorg/gradle/api/file/DirectoryProperty;
+	public abstract fun getReferenceOutputDirectory ()Lorg/gradle/api/provider/Property;
 	public abstract fun getScreenshotDir ()Lorg/gradle/api/provider/Property;
+	public abstract fun getShouldWriteReferences ()Lorg/gradle/api/provider/Property;
 	public final fun pullScreenshots ()V
 }
 
diff --git a/dropshots-gradle-plugin/build.gradle.kts b/dropshots-gradle-plugin/build.gradle.kts
index 7ce2c42..f3f547c 100644
--- a/dropshots-gradle-plugin/build.gradle.kts
+++ b/dropshots-gradle-plugin/build.gradle.kts
@@ -1,11 +1,11 @@
-import com.android.build.gradle.tasks.SourceJarTask
 import com.vanniktech.maven.publish.GradlePlugin
 import com.vanniktech.maven.publish.JavadocJar.Dokka
-import org.gradle.jvm.tasks.Jar
+import org.gradle.kotlin.dsl.provideDelegate
 import org.jetbrains.kotlin.gradle.dsl.JvmTarget
 import org.jetbrains.kotlin.gradle.dsl.KotlinVersion
-import org.jetbrains.kotlin.gradle.tasks.KotlinCompilationTask
-import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+import java.io.FileOutputStream
+import java.util.Properties
+import kotlin.apply
 
 plugins {
   `java-gradle-plugin`
@@ -15,31 +15,13 @@ plugins {
   alias(libs.plugins.binaryCompatibilityValidator)
 }
 
-buildscript {
-  repositories {
-    mavenCentral()
-    gradlePluginPortal()
-  }
-}
-
-repositories {
-  mavenCentral()
-  gradlePluginPortal()
-}
-
-sourceSets {
-  main.configure {
-    java.srcDir("src/generated/kotlin")
-  }
-}
-
 mavenPublishing {
   configure(GradlePlugin(Dokka("dokkaJavadoc")))
 }
 
-val generateVersionTask = tasks.register("generateVersion") {
+val generateVersionTask: TaskProvider<*> = tasks.register("generateVersion") {
   inputs.property("version", project.property("VERSION_NAME") as String)
-  outputs.dir(project.layout.projectDirectory.dir("src/generated/kotlin"))
+  outputs.dir(layout.buildDirectory.dir("generated/generateVersion"))
 
   doLast {
     val output = File(outputs.files.first(), "com/dropbox/dropshots/Version.kt")
@@ -52,23 +34,9 @@ val generateVersionTask = tasks.register("generateVersion") {
   }
 }
 
-tasks.withType<Jar>().configureEach {
-  dependsOn(generateVersionTask)
-}
-
-tasks.named("dokkaJavadoc").configure {
-  dependsOn(generateVersionTask)
-}
-
-tasks.named("compileKotlin").configure {
-  dependsOn(generateVersionTask)
-}
-
-tasks.withType<KotlinCompile>().configureEach {
-  compilerOptions {
-    jvmTarget.set(JvmTarget.JVM_11)
-    apiVersion.set(KotlinVersion.KOTLIN_1_8)
-    languageVersion.set(KotlinVersion.KOTLIN_1_8)
+sourceSets {
+  main.configure {
+    java.srcDir(generateVersionTask.map { it.outputs.files })
   }
 }
 
@@ -77,6 +45,11 @@ tasks.withType<JavaCompile>().configureEach {
 }
 
 kotlin {
+  compilerOptions {
+    jvmTarget.set(JvmTarget.JVM_11)
+    apiVersion.set(KotlinVersion.KOTLIN_2_0)
+    languageVersion.set(KotlinVersion.KOTLIN_2_0)
+  }
   explicitApi()
 }
 
@@ -89,20 +62,29 @@ gradlePlugin {
   }
 }
 
-// See https://github.com/slackhq/keeper/pull/11#issuecomment-579544375 for context
-val releaseMode = hasProperty("dropshots.releaseMode")
+val addTestPlugin: Configuration = configurations.create("addTestPlugin") {
+  attributes {
+    attribute(ArtifactTypeDefinition.ARTIFACT_TYPE_ATTRIBUTE, ArtifactTypeDefinition.JAR_TYPE)
+  }
+}
+
+configurations { testImplementation.get().extendsFrom(addTestPlugin) }
+
+tasks.pluginUnderTestMetadata {
+  // make sure the test can access plugins for coordination.
+  pluginClasspath.from(addTestPlugin)
+}
 dependencies {
+  addTestPlugin(gradleApi())
+  addTestPlugin(libs.android)
+  addTestPlugin(libs.kotlin.plugin)
+
+
   compileOnly(gradleApi())
+  compileOnly(libs.android)
+  compileOnly(libs.kotlin.plugin)
+
   implementation(platform(libs.kotlin.bom))
-  // Don't impose our version of KGP on consumers
-
-  if (releaseMode) {
-    compileOnly(libs.android)
-    compileOnly(libs.kotlin.plugin)
-  } else {
-    implementation(libs.android)
-    implementation(libs.kotlin.plugin)
-  }
 
   testImplementation(gradleTestKit())
   testImplementation(platform(libs.kotlin.bom))
diff --git a/dropshots-gradle-plugin/gradle b/dropshots-gradle-plugin/gradle
new file mode 120000
index 0000000..3337596
--- /dev/null
+++ b/dropshots-gradle-plugin/gradle
@@ -0,0 +1 @@
+../gradle
\ No newline at end of file
diff --git a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/ClearScreenshotsTask.kt b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/ClearScreenshotsTask.kt
index 88f2e4e..4f2e908 100644
--- a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/ClearScreenshotsTask.kt
+++ b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/ClearScreenshotsTask.kt
@@ -7,7 +7,11 @@ import org.gradle.api.tasks.Destroys
 import org.gradle.api.tasks.Input
 import org.gradle.api.tasks.TaskAction
 import org.gradle.process.ExecOperations
+import org.gradle.work.DisableCachingByDefault
 
+@DisableCachingByDefault(
+  because = "The task interacts with the connected test device, so caching is not applicable."
+)
 public abstract class ClearScreenshotsTask : DefaultTask() {
 
   @get:Input
@@ -22,6 +26,8 @@ public abstract class ClearScreenshotsTask : DefaultTask() {
   init {
     description = "Removes the test screenshots from the test device."
     group = "verification"
+
+    outputs.upToDateWhen { false }
   }
 
   @TaskAction
diff --git a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsExtension.kt b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsExtension.kt
index 662047a..9117a91 100644
--- a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsExtension.kt
+++ b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsExtension.kt
@@ -12,4 +12,12 @@ public abstract class DropshotsExtension @Inject constructor(objects: ObjectFact
    */
   public val referenceOutputDirectory: Property<String> = objects.property(String::class.java)
     .convention("src/androidTest/screenshots")
+
+  /**
+   * Whether to record screenshots on test failure. If true and recording is enabled, the
+   * screenshots from the test device will be pulled and saved to the [referenceOutputDirectory]
+   * even if the test fails.
+   */
+  public val recordOnFailure: Property<Boolean> = objects.property(Boolean::class.java)
+    .convention(true)
 }
diff --git a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsPlugin.kt b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsPlugin.kt
index ab9837c..f27871e 100644
--- a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsPlugin.kt
+++ b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/DropshotsPlugin.kt
@@ -1,130 +1,144 @@
 package com.dropbox.dropshots
 
-import com.android.build.gradle.AppExtension
-import com.android.build.gradle.LibraryExtension
-import com.android.build.gradle.TestedExtension
-import com.android.build.gradle.api.ApkVariant
-import com.android.build.gradle.internal.tasks.AndroidTestTask
-import com.android.build.gradle.internal.tasks.factory.dependsOn
-import java.util.Locale
+import com.android.build.api.dsl.AndroidSourceSet
+import com.android.build.api.dsl.ApplicationExtension
+import com.android.build.api.dsl.LibraryExtension
+import com.android.build.api.variant.AndroidComponentsExtension
+import com.android.build.api.variant.HasAndroidTest
+import com.android.build.api.variant.Variant
+import com.android.build.api.variant.impl.capitalizeFirstChar
+import kotlin.jvm.java
+import org.gradle.api.NamedDomainObjectContainer
 import org.gradle.api.Plugin
 import org.gradle.api.Project
-import org.gradle.api.tasks.Copy
+import org.gradle.api.provider.Provider
 
 private const val recordScreenshotsArg = "dropshots.record"
 
 public class DropshotsPlugin : Plugin<Project> {
   override fun apply(project: Project) {
-    val dropshotsExtension = project.extensions.create("dropshots", DropshotsExtension::class.java)
+    val dropshotsExtension: DropshotsExtension =
+      project.extensions.create("dropshots", DropshotsExtension::class.java)
 
     project.pluginManager.withPlugin("com.android.application") {
-      val extension = project.extensions.findByType(AppExtension::class.java)
-        ?: throw Exception("Failed to find Android Application extension")
-      project.configureDropshots(dropshotsExtension, extension)
+      project.extensions.getByType(ApplicationExtension::class.java).apply {
+        sourceSets.addDropshotsAssetsDir(dropshotsExtension)
+      }
+      project.configureDropshots(dropshotsExtension)
     }
 
     project.pluginManager.withPlugin("com.android.library") {
-      val extension = project.extensions.findByType(LibraryExtension::class.java)
-        ?: throw Exception("Failed to find Android Library extension")
-      project.configureDropshots(dropshotsExtension, extension)
+      project.extensions.getByType(LibraryExtension::class.java).apply {
+        sourceSets.addDropshotsAssetsDir(dropshotsExtension)
+      }
+      project.configureDropshots(dropshotsExtension)
+    }
+  }
+
+  private fun NamedDomainObjectContainer<out AndroidSourceSet>.addDropshotsAssetsDir(
+    dropshotsExtension: DropshotsExtension
+  ) {
+    named("androidTest") {
+      it.assets.directories += dropshotsExtension.referenceOutputDirectory.get()
     }
   }
 
   private fun Project.configureDropshots(
-    dropshotsExtension: DropshotsExtension,
-    extension: TestedExtension
+    dropshotsExtension: DropshotsExtension
   ) {
-    project.afterEvaluate {
-      it.dependencies.add(
+    val androidComponents = project.extensions.getByType(AndroidComponentsExtension::class.java)
+    val adbExecutablePath = androidComponents.sdkComponents.adb.map { it.asFile.path }
+
+    androidComponents.onVariants {
+      if (it is HasAndroidTest) {
+        configureVariant(it, adbExecutablePath, dropshotsExtension)
+      }
+    }
+
+    afterEvaluate {
+      dependencies.add(
         "androidTestImplementation",
         "com.dropbox.dropshots:dropshots:$VERSION"
       )
     }
+  }
 
-    //check this to have resource based on flavours
-    val androidTestSourceSet = extension.sourceSets.findByName("androidTest")
-      ?: throw Exception("Failed to find androidTest source set")
-
-    val referenceScreenshotDirectory = layout.projectDirectory.dir(dropshotsExtension.referenceOutputDirectory)
-
-    androidTestSourceSet.assets {
-      srcDirs(referenceScreenshotDirectory)
+  private fun <V> Project.configureVariant(
+    variant: V,
+    adbExecutablePath: Provider<String>,
+    dropshotsExtension: DropshotsExtension
+  ) where V : HasAndroidTest, V : Variant {
+    val androidTest = variant.androidTest ?: return
+
+    val androidTestVariantSlug = androidTest.name.capitalizeFirstChar()
+
+    val deviceScreenshotDir =
+      androidTest.applicationId.map { "/storage/emulated/0/Download/screenshots/$it" }
+    val buildScreenshotDir =
+      layout.buildDirectory.dir("test-results/dropshots/$androidTestVariantSlug")
+    val testTaskName = "connected${variant.name.capitalizeFirstChar()}AndroidTest"
+
+    val clearScreenshotsTask = tasks.register(
+      "clear${androidTestVariantSlug}Screenshots",
+      ClearScreenshotsTask::class.java,
+    ) {
+      it.adbExecutable.set(adbExecutablePath)
+      it.screenshotDir.set(deviceScreenshotDir)
     }
 
-    val adbExecutablePath = provider { extension.adbExecutable.path }
-    extension.testVariants.all { variant ->
-      val testTaskProvider = variant.connectedInstrumentTestProvider
-      val variantSlug = variant.name.replaceFirstChar { if (it.isLowerCase()) it.titlecase(Locale.getDefault()) else it.toString() }
-
-      val screenshotDir = provider {
-        val appId = if (variant.testedVariant is ApkVariant) {
-          variant.testedVariant.applicationId
-        } else {
-          variant.packageApplicationProvider.get().applicationId
-          variant.applicationId
-        }
-        "/storage/emulated/0/Download/screenshots/$appId"
-      }
+    val isRecordingScreenshots = project.objects.property(Boolean::class.java)
 
-      val clearScreenshotsTask = tasks.register(
-        "clear${variantSlug}Screenshots",
-        ClearScreenshotsTask::class.java,
-      ) {
-        it.adbExecutable.set(adbExecutablePath)
-        it.screenshotDir.set(screenshotDir)
-      }
+    val canRecordScreenshots = dropshotsExtension.recordOnFailure.map {
+      project.state.failure != null || it
+    }
 
-      val pullScreenshotsTask = tasks.register(
-        "pull${variantSlug}Screenshots",
-        PullScreenshotsTask::class.java,
-      ) {
-        it.adbExecutable.set(adbExecutablePath)
-        it.screenshotDir.set(screenshotDir)
-        it.outputDirectory.set(testTaskProvider.flatMap { (it as AndroidTestTask).resultsDir })
-        it.finalizedBy(clearScreenshotsTask)
-      }
+    val pullScreenshotsTask = tasks.register(
+      "pull${androidTestVariantSlug}Screenshots",
+      PullScreenshotsTask::class.java,
+    ) { task ->
+      task.adbExecutable.set(adbExecutablePath)
+      task.screenshotDir.set(deviceScreenshotDir)
+      task.outputDirectory.set(buildScreenshotDir)
+      task.shouldWriteReferences.set(isRecordingScreenshots.map { it && (canRecordScreenshots.get()) })
+      task.referenceOutputDirectory.set(dropshotsExtension.referenceOutputDirectory)
+    }
 
-      val recordScreenshotsTask = tasks.register(
-        "record${variantSlug}Screenshots",
-        Copy::class.java,
-      ) {
-        it.group = "verification"
-        it.description = "Updates the local reference screenshots"
-        it.from(
-          testTaskProvider.flatMap {
-            (it as AndroidTestTask).resultsDir.map { base -> "$base/reference" }
-          }
-        )
-        it.into(referenceScreenshotDirectory)
-        it.dependsOn(pullScreenshotsTask)
-        it.finalizedBy(clearScreenshotsTask)
-      }
+    val recordScreenshotsTask = tasks.register(
+      "record${androidTestVariantSlug}Screenshots",
+    ) { task ->
+      task.group = "verification"
+      task.description = "Indicates that local reference screenshots should be updated"
+      // This task being present on the task graph is used as an indicator that we should update
+      // source screenshots. The actual work is performed in pullScreenshotsTask since it is
+      // a finalizer task.
+      task.dependsOn(testTaskName)
+    }
 
-      val isRecordingScreenshots = project.objects.property(Boolean::class.java)
-      if (providers.gradleProperty(recordScreenshotsArg).isPresent) {
-        project.logger.warn("The 'dropshots.record' property has been deprecated and will " +
-          "be removed in a future version.")
-        isRecordingScreenshots.set(true)
-      }
-      project.gradle.taskGraph.whenReady { graph ->
-        isRecordingScreenshots.set(recordScreenshotsTask.map { graph.hasTask(it) })
-      }
+    if (providers.gradleProperty(recordScreenshotsArg).isPresent) {
+      project.logger.warn(
+        "The 'dropshots.record' property has been deprecated and will " +
+          "be removed in a future version."
+      )
+      isRecordingScreenshots.set(true)
+    }
+    project.gradle.taskGraph.whenReady { graph ->
+      isRecordingScreenshots.set(recordScreenshotsTask.map { graph.hasTask(it) })
+    }
 
-      val writeMarkerFileTask = tasks.register(
-        "push${variantSlug}ScreenshotMarkerFile",
-        PushFileTask::class.java,
-      ) {
-        it.onlyIf { isRecordingScreenshots.get() }
-        it.adbExecutable.set(adbExecutablePath)
-        it.fileContents.set("\n")
-        it.remotePath.set(screenshotDir.map { dir -> "$dir/.isRecordingScreenshots" })
-        it.finalizedBy(clearScreenshotsTask)
-      }
-      testTaskProvider.dependsOn(writeMarkerFileTask)
+    val writeMarkerFileTask = tasks.register(
+      "push${androidTestVariantSlug}ScreenshotMarkerFile",
+      PushFileTask::class.java,
+    ) { task ->
+      task.onlyIf { isRecordingScreenshots.get() }
+      task.adbExecutable.set(adbExecutablePath)
+      task.fileContents.set("\n")
+      task.remotePath.set(deviceScreenshotDir.map { "$it/.isRecordingScreenshots" })
+      task.dependsOn(clearScreenshotsTask)
+    }
 
-      testTaskProvider.configure {
-        it.finalizedBy(pullScreenshotsTask)
-      }
+    tasks.named { it == testTaskName }.configureEach {
+      it.finalizedBy(pullScreenshotsTask)
+      it.dependsOn(writeMarkerFileTask)
     }
   }
 }
diff --git a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PullScreenshotsTask.kt b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PullScreenshotsTask.kt
index 7e6f16b..3931b82 100644
--- a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PullScreenshotsTask.kt
+++ b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PullScreenshotsTask.kt
@@ -4,12 +4,20 @@ import javax.inject.Inject
 import java.io.ByteArrayOutputStream
 import org.gradle.api.DefaultTask
 import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.file.FileSystemOperations
 import org.gradle.api.provider.Property
 import org.gradle.api.tasks.Input
+import org.gradle.api.tasks.InputDirectory
 import org.gradle.api.tasks.OutputDirectory
+import org.gradle.api.tasks.PathSensitive
+import org.gradle.api.tasks.PathSensitivity
 import org.gradle.api.tasks.TaskAction
 import org.gradle.process.ExecOperations
+import org.gradle.work.DisableCachingByDefault
 
+@DisableCachingByDefault(
+  because = "The task interacts with the connected test device, so caching is not applicable."
+)
 public abstract class PullScreenshotsTask : DefaultTask() {
 
   @get:Input
@@ -18,12 +26,21 @@ public abstract class PullScreenshotsTask : DefaultTask() {
   @get:Input
   public abstract val screenshotDir: Property<String>
 
+  @get:Input
+  public abstract val shouldWriteReferences: Property<Boolean>
+
+  @get:Input
+  public abstract val referenceOutputDirectory: Property<String>
+
   @get:OutputDirectory
   public abstract val outputDirectory: DirectoryProperty
 
   @get:Inject
   protected abstract val execOperations: ExecOperations
 
+  @get:Inject
+  protected abstract val fileOperations: FileSystemOperations
+
   init {
     description = "Pull screenshots from the test device."
     group = "verification"
@@ -60,6 +77,13 @@ public abstract class PullScreenshotsTask : DefaultTask() {
         print(output.toString(Charsets.UTF_8))
       }
     }
+
+    if(shouldWriteReferences.get()) {
+      fileOperations.copy {
+        it.from(outputDirectory.dir("reference") )
+        it.into(referenceOutputDirectory)
+      }
+    }
   }
 }
 
diff --git a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PushFileTask.kt b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PushFileTask.kt
index 9751bb5..197abf5 100644
--- a/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PushFileTask.kt
+++ b/dropshots-gradle-plugin/src/main/kotlin/com/dropbox/dropshots/PushFileTask.kt
@@ -10,7 +10,11 @@ import org.gradle.api.provider.Property
 import org.gradle.api.tasks.Input
 import org.gradle.api.tasks.TaskAction
 import org.gradle.process.ExecOperations
+import org.gradle.work.DisableCachingByDefault
 
+@DisableCachingByDefault(
+  because = "The task interacts with the connected test device, so caching is not applicable."
+)
 public abstract class PushFileTask : DefaultTask() {
 
   @get:Input
diff --git a/dropshots-gradle-plugin/src/test/kotlin/com/dropbox/dropshots/DropshotsPluginTest.kt b/dropshots-gradle-plugin/src/test/kotlin/com/dropbox/dropshots/DropshotsPluginTest.kt
index bc952b8..97353cd 100644
--- a/dropshots-gradle-plugin/src/test/kotlin/com/dropbox/dropshots/DropshotsPluginTest.kt
+++ b/dropshots-gradle-plugin/src/test/kotlin/com/dropbox/dropshots/DropshotsPluginTest.kt
@@ -7,15 +7,29 @@ import org.gradle.testkit.runner.TaskOutcome
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
 import org.junit.rules.TemporaryFolder
 
-class DropshotsPluginTest {
+@RunWith(Parameterized::class)
+class DropshotsPluginTest(private val agpVersion: String) {
 
   @get:Rule val tmpFolder = TemporaryFolder()
 
   private lateinit var buildFile: File
   private lateinit var gradleRunner: GradleRunner
 
+  companion object {
+    @JvmStatic
+    @Parameterized.Parameters(name = "AGP {0}")
+    fun agpVersions(): Collection<Array<Any>> {
+      return listOf(
+        arrayOf("8.7.2"),
+        arrayOf("9.1.0"),
+      )
+    }
+  }
+
   @Before
   fun setup() {
     val localMavenRepo = File("../build/localMaven").absolutePath
@@ -59,7 +73,7 @@ class DropshotsPluginTest {
       // language=groovy
       """
         plugins {
-          alias(libs.plugins.android.library)
+          id("com.android.library") version "$agpVersion"
           id("com.dropbox.dropshots")
         }
 
@@ -97,7 +111,7 @@ class DropshotsPluginTest {
         """
           plugins {
             id("com.dropbox.dropshots")
-            alias(libs.plugins.android.library)
+            id("com.android.library") version "$agpVersion"
           }
 
           android {
@@ -121,8 +135,9 @@ class DropshotsPluginTest {
   @Test
   fun `executes marker file push only when record task is run`() {
     val result = gradleRunner
-      .withArguments("recordDebugAndroidTestScreenshots")
+      .withArguments("recordDebugAndroidTestScreenshots", "--stacktrace")
       .build()
+    println(result.output)
     with(result.task(":pushDebugAndroidTestScreenshotMarkerFile")) {
       assertThat(this).isNotNull()
       assertThat(this!!.outcome).isEqualTo(TaskOutcome.SUCCESS)
@@ -148,7 +163,7 @@ class DropshotsPluginTest {
         """
           plugins {
             id("com.dropbox.dropshots")
-            alias(libs.plugins.android.library)
+            id("com.android.library") version "$agpVersion"
           }
 
           dropshots {
@@ -190,7 +205,7 @@ class DropshotsPluginTest {
         """
           plugins {
             id("com.dropbox.dropshots")
-            alias(libs.plugins.android.library)
+            id("com.android.library") version "$agpVersion"
           }
 
           android {
diff --git a/dropshots/build.gradle.kts b/dropshots/build.gradle.kts
index 0e59338..da58dfe 100644
--- a/dropshots/build.gradle.kts
+++ b/dropshots/build.gradle.kts
@@ -1,10 +1,8 @@
 import com.vanniktech.maven.publish.AndroidSingleVariantLibrary
-import java.io.ByteArrayOutputStream
-import java.util.Locale
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget
 
 plugins {
   alias(libs.plugins.android.library)
-  alias(libs.plugins.kotlin.android)
   alias(libs.plugins.dokka)
   alias(libs.plugins.mavenPublish)
   alias(libs.plugins.binaryCompatibilityValidator)
@@ -25,13 +23,13 @@ android {
     sourceCompatibility = JavaVersion.VERSION_1_8
     targetCompatibility = JavaVersion.VERSION_1_8
   }
-  kotlinOptions {
-    jvmTarget = "1.8"
-  }
 }
 
 kotlin {
   explicitApi()
+  compilerOptions {
+    jvmTarget = JvmTarget.JVM_1_8
+  }
 }
 
 dependencies {
@@ -41,12 +39,11 @@ dependencies {
   implementation(libs.androidx.test.runner)
   implementation(libs.androidx.test.rules)
 
-  debugImplementation(libs.androidx.fragment)
-
   testImplementation(platform(libs.kotlin.bom))
   testImplementation(libs.junit)
   testImplementation(libs.kotlin.test)
 
+  androidTestImplementation(libs.androidx.fragment)
   androidTestImplementation(libs.junit)
   androidTestImplementation(libs.androidx.test.core)
   androidTestImplementation(libs.androidx.test.ext.junit)
@@ -62,125 +59,3 @@ mavenPublishing {
     publishJavadocJar = true,
   ))
 }
-
-val adbExecutablePath = provider { android.adbExecutable.path }
-android.testVariants.all {
-  val screenshotDir = "/storage/emulated/0/Download/screenshots/com.dropbox.dropshots.test"
-  val connectedAndroidTest = connectedInstrumentTestProvider
-  val variantSlug = name.replaceFirstChar { if (it.isLowerCase()) it.titlecase(Locale.getDefault()) else it.toString() }
-
-  val recordScreenshotsTask = tasks.register("record${variantSlug}Screenshots")
-  val isRecordingScreenshots = project.objects.property(Boolean::class.java)
-  project.gradle.taskGraph.whenReady {
-    isRecordingScreenshots.set(recordScreenshotsTask.map { hasTask(it) })
-  }
-
-  val pushMarkerFileTask = tasks.register("push${variantSlug}ScreenshotMarkerFile") {
-    description = "Push screenshot marker file to test device."
-    group = "verification"
-    outputs.upToDateWhen { false }
-    onlyIf { isRecordingScreenshots.get() }
-
-    doLast {
-      val adb = adbExecutablePath.get()
-      project.exec {
-        executable = adb
-        args = listOf("shell", "mkdir", "-p", screenshotDir)
-      }
-      project.exec {
-        executable = adb
-        args = listOf("shell", "touch", "$screenshotDir/.isRecordingScreenshots")
-      }
-    }
-  }
-
-  val setupEmulatorTask = tasks.register("setup${variantSlug}ScreenshotEmulator") {
-    description = "Configures the test device for screenshots."
-    group = "verification"
-    doLast {
-      val adb = adbExecutablePath.get()
-      fun adbCommand(cmd: String): ExecResult {
-        return project.exec {
-          executable = adb
-          args = cmd.split(" ")
-        }
-      }
-
-      adbCommand("root")
-      adbCommand("wait-for-device")
-      adbCommand("shell settings put global sysui_demo_allowed 1")
-      adbCommand("shell am broadcast -a com.android.systemui.demo -e command enter")
-        .assertNormalExitValue()
-      adbCommand("shell am broadcast -a com.android.systemui.demo -e command clock -e hhmm 1234")
-      adbCommand("shell am broadcast -a com.android.systemui.demo -e command battery -e plugged false")
-      adbCommand("shell am broadcast -a com.android.systemui.demo -e command battery -e level 100")
-      adbCommand("shell am broadcast -a com.android.systemui.demo -e command network -e wifi show -e level 4")
-      adbCommand("shell am broadcast -a com.android.systemui.demo -e command network -e mobile show -e datatype none -e level 4")
-      adbCommand("shell am broadcast -a com.android.systemui.demo -e command notifications -e visible false")
-    }
-  }
-  val restoreEmulatorTask = tasks.register("restore${variantSlug}ScreenshotEmulator") {
-    description = "Restores the test device from screenshot mode."
-    group = "verification"
-    doLast {
-      project.exec {
-        executable = adbExecutablePath.get()
-        args = "shell am broadcast -a com.android.systemui.demo -e command exit".split(" ")
-      }
-    }
-  }
-
-  val pullScreenshotsTask = tasks.register("pull${variantSlug}Screenshots") {
-    description = "Pull screenshots from the test device."
-    group = "verification"
-    outputs.dir(project.layout.buildDirectory.dir("reports/androidTests/dropshots"))
-    outputs.upToDateWhen { false }
-
-    doLast {
-      val outputDir = outputs.files.singleFile
-      outputDir.mkdirs()
-
-      val adb = adbExecutablePath.get()
-      val checkResult = project.exec {
-        executable = adb
-        args = listOf("shell", "test", "-d", screenshotDir)
-        isIgnoreExitValue = true
-      }
-
-      if (checkResult.exitValue == 0) {
-        val output = ByteArrayOutputStream()
-        val pullResult = project.exec {
-          executable = adb
-          args = listOf("pull", "$screenshotDir/.", outputDir.path)
-          standardOutput = output
-          isIgnoreExitValue = true
-        }
-
-        if (pullResult.exitValue == 0) {
-          val fileCount = """^${screenshotDir.replace(".", "\\.")}/?\./: ([0-9]*) files pulled,.*$""".toRegex()
-          val matchResult = fileCount.find(output.toString(Charsets.UTF_8))
-          if (matchResult != null && matchResult.groups.size > 1) {
-            println("${matchResult.groupValues[1]} screenshots saved at ${outputDir.path}")
-          } else {
-            println("Unknown result executing adb: $adb pull $screenshotDir/. ${outputDir.path}")
-            print(output.toString(Charsets.UTF_8))
-          }
-        } else {
-          println("Failed to pull screenshots.")
-        }
-
-        project.exec {
-          executable = adb
-          args = listOf("shell", "rm", "-r", "/storage/emulated/0/Download/screenshots")
-        }
-      }
-    }
-  }
-
-  connectedAndroidTest.configure {
-    dependsOn(pushMarkerFileTask)
-    finalizedBy(pullScreenshotsTask)
-    dependsOn(setupEmulatorTask)
-    finalizedBy(restoreEmulatorTask)
-  }
-}
diff --git a/dropshots/src/androidTest/assets/MatchesFullScreenshot.png b/dropshots/src/androidTest/assets/MatchesFullScreenshot.png
deleted file mode 100644
index d7c93da..0000000
Binary files a/dropshots/src/androidTest/assets/MatchesFullScreenshot.png and /dev/null differ
diff --git a/dropshots/src/main/java/com/dropbox/dropshots/Dropshots.kt b/dropshots/src/main/java/com/dropbox/dropshots/Dropshots.kt
index fea5b15..81c7de6 100644
--- a/dropshots/src/main/java/com/dropbox/dropshots/Dropshots.kt
+++ b/dropshots/src/main/java/com/dropbox/dropshots/Dropshots.kt
@@ -21,12 +21,30 @@ import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
 
-public class Dropshots internal constructor(
-  private val rootScreenshotDirectory: File,
-  private val filenameFunc: (String, String) -> String,
-  private val recordScreenshots: Boolean,
-  private val imageComparator: ImageComparator,
-  private val resultValidator: ResultValidator,
+
+public class Dropshots @JvmOverloads public constructor(
+  /**
+   * Function to create a filename from the class name and snapshot name (i.e. the name provided when taking
+   * the snapshot). Default behavior will use the calling function name.
+   */
+  private val filenameFunc: (String, String) -> String = defaultFilenameFunc,
+  /**
+   * Indicates whether new reference screenshots should be recorded. Otherwise Dropshots performs
+   * validation of test screenshots against reference screenshots.
+   */
+  public val recordScreenshots: Boolean = isRecordingScreenshots(defaultRootScreenshotDirectory()),
+  /**
+   * The `ImageComparator` used to compare test and reference screenshots.
+   */
+  private val imageComparator: ImageComparator = SimpleImageComparator(maxDistance = 0.004f),
+  /**
+   * The `ResultValidator` used to validate the comparison results.
+   */
+  private val resultValidator: ResultValidator = CountValidator(0),
+  /**
+   * Device directory where screenshots will be written.
+   */
+  public val rootScreenshotDirectory: File = defaultRootScreenshotDirectory(),
 ) : TestRule {
   private val context = InstrumentationRegistry.getInstrumentation().context
   private var fqName: String = ""
@@ -36,28 +54,6 @@ public class Dropshots internal constructor(
 
   private val snapshotName: String get() = testName
 
-  @JvmOverloads
-  public constructor(
-    /**
-     * Function to create a filename from the class name and snapshot name (i.e. the name provided when taking
-     * the snapshot). Default behavior will use the calling function name.
-     */
-    filenameFunc: (String, String) -> String = defaultFilenameFunc,
-    /**
-     * Indicates whether new reference screenshots should be recorded. Otherwise Dropshots performs
-     * validation of test screenshots against reference screenshots.
-     */
-    recordScreenshots: Boolean = isRecordingScreenshots(defaultRootScreenshotDirectory()),
-    /**
-     * The `ImageComparator` used to compare test and reference screenshots.
-     */
-    imageComparator: ImageComparator = SimpleImageComparator(maxDistance = 0.004f),
-    /**
-     * The `ResultValidator` used to validate the comparison results.
-     */
-    resultValidator: ResultValidator = CountValidator(0),
-  ): this(defaultRootScreenshotDirectory(), filenameFunc, recordScreenshots, imageComparator, resultValidator)
-
   override fun apply(base: Statement, description: Description): Statement {
     fqName = description.className
     packageName = fqName.substringBeforeLast('.', missingDelimiterValue = "")
@@ -221,7 +217,7 @@ public class Dropshots internal constructor(
 
     val file = File(dir, "${name.replace(" ", "_")}.png")
     if (file.exists()){
-      throw IllegalStateException("Unable to create screenshot, file already exists. Please " +
+      throw IllegalStateException("Unable to create screenshot, ${file.path} already exists. Please " +
         "specify name param with something more specific when calling assertSnapshot function")
     }
 
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 0e2c2bd..393c89b 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -1,6 +1,6 @@
 [versions]
-kotlin = "2.0.21"
-agp = "8.7.2"
+kotlin = "2.2.21"
+agp = "9.1.0"
 
 [libraries]
 android = { module = "com.android.tools.build:gradle", version.ref = "agp" }
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index 5c40527..e1f0f9a 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,6 +1,6 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-8.10.2-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-9.4.0-bin.zip
 networkTimeout=10000
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
diff --git a/settings.gradle.kts b/settings.gradle.kts
index 4cab8bb..f12c85e 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -6,6 +6,14 @@ pluginManagement {
   }
 }
 
+dependencyResolutionManagement {
+  repositories {
+    mavenCentral()
+    google()
+    gradlePluginPortal()
+  }
+}
+
 rootProject.name = "dropshots-root"
 
 include(":dropshots-gradle-plugin")
diff --git a/tests/build.gradle.kts b/tests/build.gradle.kts
new file mode 100644
index 0000000..bc1d231
--- /dev/null
+++ b/tests/build.gradle.kts
@@ -0,0 +1,30 @@
+plugins {
+  alias(libs.plugins.android.library)
+  id("com.dropbox.dropshots")
+}
+
+dropshots {
+  referenceOutputDirectory = "src/androidTest/assets"
+}
+
+android {
+  namespace = "com.dropbox.dropshots.integrationtests"
+  compileSdk = 34
+
+  defaultConfig {
+    minSdk = 21
+
+    testInstrumentationRunner = "androidx.test.runner.AndroidJUnitRunner"
+  }
+}
+
+dependencies {
+  implementation(libs.androidx.fragment)
+
+  androidTestImplementation(libs.junit)
+  androidTestImplementation(libs.androidx.test.core)
+  androidTestImplementation(libs.androidx.test.ext.junit)
+  androidTestImplementation(libs.androidx.test.rules)
+  androidTestImplementation(libs.androidx.test.runner)
+  androidTestImplementation(libs.androidx.test.uiautomator)
+}
diff --git a/tests/gradle b/tests/gradle
new file mode 120000
index 0000000..3337596
--- /dev/null
+++ b/tests/gradle
@@ -0,0 +1 @@
+../gradle
\ No newline at end of file
diff --git a/tests/gradlew b/tests/gradlew
new file mode 120000
index 0000000..502f5a2
--- /dev/null
+++ b/tests/gradlew
@@ -0,0 +1 @@
+../gradlew
\ No newline at end of file
diff --git a/tests/settings.gradle.kts b/tests/settings.gradle.kts
new file mode 100644
index 0000000..9ee27e7
--- /dev/null
+++ b/tests/settings.gradle.kts
@@ -0,0 +1,21 @@
+pluginManagement {
+  repositories {
+    gradlePluginPortal()
+    mavenCentral()
+    google()
+  }
+  includeBuild("..")
+}
+
+dependencyResolutionManagement {
+  repositories {
+    mavenCentral()
+    google()
+  }
+}
+
+includeBuild("..") {
+  dependencySubstitution {
+    substitute(module("com.dropbox:dropshots")).using(project(":dropshots"))
+  }
+}
diff --git a/dropshots/src/androidTest/assets/MatchesActivityScreenshot.png b/tests/src/androidTest/assets/MatchesActivityScreenshot.png
similarity index 100%
rename from dropshots/src/androidTest/assets/MatchesActivityScreenshot.png
rename to tests/src/androidTest/assets/MatchesActivityScreenshot.png
diff --git a/tests/src/androidTest/assets/MatchesFullScreenshot.png b/tests/src/androidTest/assets/MatchesFullScreenshot.png
new file mode 100644
index 0000000..25fd285
Binary files /dev/null and b/tests/src/androidTest/assets/MatchesFullScreenshot.png differ
diff --git a/dropshots/src/androidTest/assets/MatchesViewScreenshot.png b/tests/src/androidTest/assets/MatchesViewScreenshot.png
similarity index 100%
rename from dropshots/src/androidTest/assets/MatchesViewScreenshot.png
rename to tests/src/androidTest/assets/MatchesViewScreenshot.png
diff --git a/dropshots/src/androidTest/assets/static/50x50.png b/tests/src/androidTest/assets/static/50x50.png
similarity index 100%
rename from dropshots/src/androidTest/assets/static/50x50.png
rename to tests/src/androidTest/assets/static/50x50.png
diff --git a/dropshots/src/androidTest/assets/static/MatchesViewScreenshotBad.png b/tests/src/androidTest/assets/static/MatchesViewScreenshotBad.png
similarity index 100%
rename from dropshots/src/androidTest/assets/static/MatchesViewScreenshotBad.png
rename to tests/src/androidTest/assets/static/MatchesViewScreenshotBad.png
diff --git a/dropshots/src/androidTest/assets/static/MatchesViewScreenshotBadSize.png b/tests/src/androidTest/assets/static/MatchesViewScreenshotBadSize.png
similarity index 100%
rename from dropshots/src/androidTest/assets/static/MatchesViewScreenshotBadSize.png
rename to tests/src/androidTest/assets/static/MatchesViewScreenshotBadSize.png
diff --git a/dropshots/src/androidTest/kotlin/com/dropbox/dropshots/CustomImageComparatorTest.kt b/tests/src/androidTest/kotlin/com/dropbox/dropshots/CustomImageComparatorTest.kt
similarity index 96%
rename from dropshots/src/androidTest/kotlin/com/dropbox/dropshots/CustomImageComparatorTest.kt
rename to tests/src/androidTest/kotlin/com/dropbox/dropshots/CustomImageComparatorTest.kt
index 905aca7..69844ce 100644
--- a/dropshots/src/androidTest/kotlin/com/dropbox/dropshots/CustomImageComparatorTest.kt
+++ b/tests/src/androidTest/kotlin/com/dropbox/dropshots/CustomImageComparatorTest.kt
@@ -26,8 +26,6 @@ class CustomImageComparatorTest {
   @get:Rule
   val dropshots = Dropshots(
     rootScreenshotDirectory = imageDirectory,
-    filenameFunc = defaultFilenameFunc,
-    recordScreenshots = false,
     imageComparator = comparator,
     resultValidator = CountValidator(0),
   )
@@ -46,6 +44,7 @@ class CustomImageComparatorTest {
     imageDirectory.deleteRecursively()
   }
 
+
   @Test
   fun imageComparatorIsConfigurable() {
     val calls = mutableListOf<Triple<Image, Image, Mask?>>()
diff --git a/dropshots/src/androidTest/kotlin/com/dropbox/dropshots/DropshotsTest.kt b/tests/src/androidTest/kotlin/com/dropbox/dropshots/DropshotsTest.kt
similarity index 96%
rename from dropshots/src/androidTest/kotlin/com/dropbox/dropshots/DropshotsTest.kt
rename to tests/src/androidTest/kotlin/com/dropbox/dropshots/DropshotsTest.kt
index 6ce91ff..4638ab6 100644
--- a/dropshots/src/androidTest/kotlin/com/dropbox/dropshots/DropshotsTest.kt
+++ b/tests/src/androidTest/kotlin/com/dropbox/dropshots/DropshotsTest.kt
@@ -4,7 +4,6 @@ import android.graphics.Bitmap
 import android.graphics.Canvas
 import android.graphics.Color
 import android.os.Environment
-import android.view.View
 import androidx.test.ext.junit.rules.ActivityScenarioRule
 import com.dropbox.differ.SimpleImageComparator
 import java.io.File
@@ -25,14 +24,12 @@ class DropshotsTest {
 
   private val fakeValidator = FakeResultValidator()
   private var filenameFunc: (String, String) -> String = { _, funcName -> funcName }
-  private val isRecordingScreenshots = isRecordingScreenshots(defaultRootScreenshotDirectory())
   private lateinit var imageDirectory: File
 
   @get:Rule val testName = TestName()
   @get:Rule
   val dropshots = Dropshots(
     filenameFunc = filenameFunc,
-    recordScreenshots = isRecordingScreenshots,
     resultValidator = fakeValidator,
     imageComparator = SimpleImageComparator(
       maxDistance = 0.004f,
@@ -43,8 +40,7 @@ class DropshotsTest {
 
   @Before
   fun setup() {
-    imageDirectory =
-      File(
+    imageDirectory = File(
         Environment.getExternalStoragePublicDirectory(Environment.DIRECTORY_DOWNLOADS),
         "screenshots/test-${testName.methodName}",
       )
@@ -58,6 +54,7 @@ class DropshotsTest {
     }
   }
 
+
   @Test
   fun testMatchesFullScreenshot() {
     activityScenarioRule.scenario.onActivity {
@@ -76,7 +73,7 @@ class DropshotsTest {
   fun testMatchesViewScreenshot() {
     activityScenarioRule.scenario.onActivity {
       dropshots.assertSnapshot(
-        it.findViewById<View>(android.R.id.content),
+        it.findViewById(android.R.id.content),
         name = "MatchesViewScreenshot"
       )
     }
diff --git a/dropshots/src/androidTest/kotlin/com/dropbox/dropshots/FakeResultValidator.kt b/tests/src/androidTest/kotlin/com/dropbox/dropshots/FakeResultValidator.kt
similarity index 100%
rename from dropshots/src/androidTest/kotlin/com/dropbox/dropshots/FakeResultValidator.kt
rename to tests/src/androidTest/kotlin/com/dropbox/dropshots/FakeResultValidator.kt
diff --git a/dropshots/src/androidTest/kotlin/com/dropbox/dropshots/ScreenshotTestFragment.kt b/tests/src/androidTest/kotlin/com/dropbox/dropshots/ScreenshotTestFragment.kt
similarity index 100%
rename from dropshots/src/androidTest/kotlin/com/dropbox/dropshots/ScreenshotTestFragment.kt
rename to tests/src/androidTest/kotlin/com/dropbox/dropshots/ScreenshotTestFragment.kt
diff --git a/dropshots/src/debug/AndroidManifest.xml b/tests/src/main/AndroidManifest.xml
similarity index 100%
rename from dropshots/src/debug/AndroidManifest.xml
rename to tests/src/main/AndroidManifest.xml
diff --git a/dropshots/src/debug/kotlin/com/dropbox/dropshots/TestActivity.kt b/tests/src/main/kotlin/com/dropbox/dropshots/TestActivity.kt
similarity index 91%
rename from dropshots/src/debug/kotlin/com/dropbox/dropshots/TestActivity.kt
rename to tests/src/main/kotlin/com/dropbox/dropshots/TestActivity.kt
index 8f537b2..ce88419 100644
--- a/dropshots/src/debug/kotlin/com/dropbox/dropshots/TestActivity.kt
+++ b/tests/src/main/kotlin/com/dropbox/dropshots/TestActivity.kt
@@ -6,7 +6,7 @@ import android.view.View
 import android.view.ViewGroup
 import android.widget.LinearLayout
 
-public class TestActivity : androidx.fragment.app.FragmentActivity()  {
+class TestActivity : androidx.fragment.app.FragmentActivity()  {
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
     setContentView(