diff --git a/build-logic/convention/src/main/kotlin/com/google/samples/apps/nowinandroid/NiaFlavor.kt b/build-logic/convention/src/main/kotlin/com/google/samples/apps/nowinandroid/NiaFlavor.kt
index f57e634cc..0a74748aa 100644
--- a/build-logic/convention/src/main/kotlin/com/google/samples/apps/nowinandroid/NiaFlavor.kt
+++ b/build-logic/convention/src/main/kotlin/com/google/samples/apps/nowinandroid/NiaFlavor.kt
@@ -1,7 +1,6 @@
 package com.google.samples.apps.nowinandroid
 
 import com.android.build.api.dsl.ApplicationExtension
-import com.android.build.api.dsl.ApplicationProductFlavor
 import com.android.build.api.dsl.CommonExtension
 import com.android.build.api.dsl.ProductFlavor
 
@@ -19,11 +18,25 @@ enum class NiaFlavor(val dimension: FlavorDimension, val applicationIdSuffix: St
     prod(FlavorDimension.contentType),
 }
 
-fun configureFlavors(
-    commonExtension: CommonExtension<*, *, *, *, *, *>,
+fun <E : CommonExtension<*, *, *, F, *, *>, F : ProductFlavor> configureFlavors(
+    commonExtension: E,
     flavorConfigurationBlock: ProductFlavor.(flavor: NiaFlavor) -> Unit = {},
 ) {
-    commonExtension.apply {
+    when (commonExtension) {
+        is ApplicationExtension -> commonExtension.configureFlavors { nia ->
+            flavorConfigurationBlock(this, nia)
+            this.applicationIdSuffix = nia.applicationIdSuffix
+        }
+
+        else -> commonExtension.configureFlavors(flavorConfigurationBlock)
+    }
+}
+
+@JvmName("configureFlavorsImpl")
+private inline fun <E : CommonExtension<*, *, *, F, *, *>, F : ProductFlavor> E.configureFlavors(
+    crossinline flavorConfigurationBlock: F.(flavor: NiaFlavor) -> Unit = {},
+) {
+    apply {
         FlavorDimension.values().forEach { flavorDimension ->
             flavorDimensions += flavorDimension.name
         }
@@ -33,11 +46,6 @@ fun configureFlavors(
                 register(niaFlavor.name) {
                     dimension = niaFlavor.dimension.name
                     flavorConfigurationBlock(this, niaFlavor)
-                    if (this@apply is ApplicationExtension && this is ApplicationProductFlavor) {
-                        if (niaFlavor.applicationIdSuffix != null) {
-                            applicationIdSuffix = niaFlavor.applicationIdSuffix
-                        }
-                    }
                 }
             }
         }