2 files changed, 3 insertions, 8 deletions
diff --git a/android/test/baselineprofile/build.gradle.kts b/android/test/baselineprofile/build.gradle.kts
index 6477b5adca..c7216bd9c5 100644
--- a/android/test/baselineprofile/build.gradle.kts
+++ b/android/test/baselineprofile/build.gradle.kts
@@ -86,6 +86,7 @@ dependencies {
 }
 
 androidComponents {
+    beforeVariants { it.enable = Variant(it.buildType, it.productFlavors).matches(baselineFilter) }
     onVariants { v ->
         val artifactsLoader = v.artifacts.getBuiltArtifactsLoader()
         v.instrumentationRunnerArguments.put(
diff --git a/android/test/e2e/build.gradle.kts b/android/test/e2e/build.gradle.kts
index 5fe61b878f..d4ee050a06 100644
--- a/android/test/e2e/build.gradle.kts
+++ b/android/test/e2e/build.gradle.kts
@@ -1,4 +1,3 @@
-import org.gradle.internal.extensions.stdlib.capitalized
 import org.jetbrains.kotlin.gradle.dsl.JvmTarget
 
 plugins {
@@ -94,13 +93,8 @@ android {
 androidComponents {
     beforeVariants { variantBuilder ->
         variantBuilder.enable =
-            variantBuilder.let { currentVariant ->
-                val enabledVariants =
-                    enabledE2eVariantTriples.map { (billing, infra, buildType) ->
-                        billing + infra.capitalized() + buildType.capitalized()
-                    }
-                enabledVariants.contains(currentVariant.name)
-            }
+            Variant(variantBuilder.buildType, variantBuilder.productFlavors)
+                .matchesAny(ossProdDebug, playStagemoleDebug)
     }
 }