diff --git a/fladle-plugin/src/main/java/com/osacky/flank/gradle/FladlePluginDelegate.kt b/fladle-plugin/src/main/java/com/osacky/flank/gradle/FladlePluginDelegate.kt
index 19759cbb..bbb173be 100644
--- a/fladle-plugin/src/main/java/com/osacky/flank/gradle/FladlePluginDelegate.kt
+++ b/fladle-plugin/src/main/java/com/osacky/flank/gradle/FladlePluginDelegate.kt
@@ -2,6 +2,7 @@ package com.osacky.flank.gradle
 
 import com.android.build.gradle.AppExtension
 import com.android.build.gradle.TestedExtension
+import com.android.build.gradle.api.BaseVariant
 import com.android.builder.model.TestOptions
 import com.osacky.flank.gradle.validation.checkForExclusionUsage
 import com.osacky.flank.gradle.validation.validateOptionsUsed
@@ -108,8 +109,9 @@ class FladlePluginDelegate {
       }
       dependsOn(writeConfigProps)
       if (config.dependOnAssemble.isPresent && config.dependOnAssemble.get()) {
-        project.extensions.findByType(TestedExtension::class.java)?.let { testedExtension ->
-          testedExtension.testVariants.configureEach {
+        val testedExtension = requireNotNull(project.extensions.findByType(TestedExtension::class.java)) { "Could not find TestedExtension in ${project.name}" }
+        testedExtension.testVariants.configureEach {
+          if (base.isExpectedVariant(testedVariant)) {
             if (testedVariant.assembleProvider.isPresent) {
               dependsOn(testedVariant.assembleProvider)
             }
@@ -143,7 +145,7 @@ class FladlePluginDelegate {
       val appVariant = testedVariant
       outputs.configureEach test@{
         appVariant.outputs.configureEach app@{
-          if (!extension.variant.isPresent || (extension.variant.isPresent && extension.variant.get() == appVariant.name)) {
+          if (extension.isExpectedVariant(appVariant)) {
 
             if (!extension.debugApk.isPresent) {
               // Don't set debug apk if not already set. #172
@@ -161,6 +163,10 @@ class FladlePluginDelegate {
     }
   }
 
+  private fun FlankGradleExtension.isExpectedVariant(
+    appVariant: BaseVariant
+  ) = !variant.isPresent || (variant.isPresent && variant.get() == appVariant.name)
+
   private val Project.fladleConfig: Configuration
     get() = configurations.getByName(FLADLE_CONFIG)
 
diff --git a/fladle-plugin/src/test/java/com/osacky/flank/gradle/integration/FlankGradlePluginIntegrationTest.kt b/fladle-plugin/src/test/java/com/osacky/flank/gradle/integration/FlankGradlePluginIntegrationTest.kt
index 2b0deb5a..af3db500 100644
--- a/fladle-plugin/src/test/java/com/osacky/flank/gradle/integration/FlankGradlePluginIntegrationTest.kt
+++ b/fladle-plugin/src/test/java/com/osacky/flank/gradle/integration/FlankGradlePluginIntegrationTest.kt
@@ -128,64 +128,20 @@ class FlankGradlePluginIntegrationTest {
     assertThat(result.output).contains("debugApk must be specified")
   }
 
-  private fun setUpDependOnAssemble(dependsOnAssemble: Boolean): BuildResult {
-    writeBuildGradle(
-      """plugins {
-          id "com.osacky.fladle"
-          id "com.android.application"
-         }
-         repositories {
-              google()
-              mavenCentral()
-          }
-         android {
-             compileSdkVersion 29
-             defaultConfig {
-                 applicationId "com.osacky.flank.gradle.sample"
-                 minSdkVersion 23
-                 targetSdkVersion 29
-                 versionCode 1
-                 versionName "1.0"
-                 testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner"
-             }
-             testOptions {
-                 execution 'ANDROIDX_TEST_ORCHESTRATOR'
-              }
-         }
-         fladle {
-           serviceAccountCredentials = project.layout.projectDirectory.file("foo")
-           dependOnAssemble = $dependsOnAssemble
-         }
-      """.trimIndent()
-    )
-    testProjectRoot.newFile("foo").writeText("{}")
-    testProjectRoot.newFolder("src/main")
-    testProjectRoot.newFile("src/main/AndroidManifest.xml").writeText(
-      """
-        <?xml version="1.0" encoding="utf-8"?>
-        <manifest package="com.osacky.flank.gradle.sample" xmlns:android="http://schemas.android.com/apk/res/android" />
-      """.trimIndent()
-    )
-    val result = GradleRunner.create()
-      .withProjectDir(testProjectRoot.root)
-      .withPluginClasspath()
-      .withArguments("runFlank", "--dry-run")
-      .build()
-
-    return result
-  }
-
   @Test
-  fun testWithDependOnAssemble() {
+  fun testWithDependOnAssembleAndFlavors() {
     val result = setUpDependOnAssemble(true)
-    assertThat(result.output).contains(":assembleDebug")
-    assertThat(result.output).contains(":assembleDebugAndroidTest")
-    assertThat(result.output).doesNotContain(":assembleRelease")
+    assertThat(result.output).contains(":assembleChocolateDebug")
+    assertThat(result.output).contains(":assembleChocolateDebugAndroidTest")
+    assertThat(result.output).doesNotContain(":assembleChocolateRelease")
+    assertThat(result.output).doesNotContain(":assembleVanilla")
   }
 
   @Test
-  fun testWithOutDependOnAssemble() {
+  fun testWithOutDependOnAssembleAndFlavors() {
     val result = setUpDependOnAssemble(false)
+    assertThat(result.output).doesNotMatch(":assemble.*")
+    assertThat(result.output).doesNotContain(":assembleChocolateDebug")
     assertThat(result.output).doesNotContain(":assembleDebug")
     assertThat(result.output).doesNotContain(":assembleRelease")
     assertThat(result.output).doesNotContain(":assembleDebugAndroidTest")
@@ -229,9 +185,7 @@ class FlankGradlePluginIntegrationTest {
       """.trimIndent()
     )
     testProjectRoot.newFile("foo").writeText("{}")
-    val result = GradleRunner.create()
-      .withProjectDir(testProjectRoot.root)
-      .withPluginClasspath()
+    val result = testProjectRoot.gradleRunner()
       .withGradleVersion(minSupportGradleVersion)
       .withArguments("printYml")
       .buildAndFail()
@@ -257,12 +211,56 @@ class FlankGradlePluginIntegrationTest {
          """.trimMargin()
     )
     testProjectRoot.newFile("foo").writeText("{}")
-    val result = GradleRunner.create()
-      .withProjectDir(testProjectRoot.root)
-      .withPluginClasspath()
-      .withGradleVersion("7.0-milestone-3")
+    val result = testProjectRoot.gradleRunner()
+      .withGradleVersion("7.0-rc-1")
       .withArguments("printYmlFooConfig")
       .build()
     assertThat(result.task(":printYmlFooConfig")!!.outcome).isEqualTo(TaskOutcome.SUCCESS)
   }
+
+  private fun setUpDependOnAssemble(dependsOnAssemble: Boolean): BuildResult {
+    testProjectRoot.setupFixture("android-project")
+    writeBuildGradle(
+      """plugins {
+          id "com.osacky.fladle"
+          id "com.android.application"
+         }
+         repositories {
+              google()
+              mavenCentral()
+          }
+         android {
+             compileSdkVersion 29
+             defaultConfig {
+                 applicationId "com.osacky.flank.gradle.sample"
+                 minSdkVersion 23
+                 targetSdkVersion 29
+                 versionCode 1
+                 versionName "1.0"
+                 testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner"
+             }
+             testOptions {
+                 execution 'ANDROIDX_TEST_ORCHESTRATOR'
+             }
+             flavorDimensions "flavor"
+             productFlavors {
+                 chocolate {
+                     dimension "flavor"
+                 }
+                 vanilla {
+                     dimension "flavor"
+                 }
+             }
+         }
+         fladle {
+           serviceAccountCredentials = project.layout.projectDirectory.file("foo")
+           dependOnAssemble = $dependsOnAssemble
+           variant = "chocolateDebug"
+         }
+      """.trimIndent()
+    )
+    return testProjectRoot.gradleRunner()
+      .withArguments("runFlank", "--dry-run")
+      .build()
+  }
 }