diff --git a/generators/build.gradle.kts b/generators/build.gradle.kts
index 900fd0dbfbd..b924f652fb9 100644
--- a/generators/build.gradle.kts
+++ b/generators/build.gradle.kts
@@ -8,7 +8,7 @@ sourceSets {
     "test" { projectDefault() }
 }
 
-fun extraSourceSet(name: String, extendMain: Boolean = true): Pair<SourceSet, Configuration> {
+fun extraSourceSet(name: String, extendMain: Boolean = true, jpsKind: String? = null): Pair<SourceSet, Configuration> {
     val sourceSet = sourceSets.create(name) {
         java.srcDir(name)
     }
@@ -18,6 +18,10 @@ fun extraSourceSet(name: String, extendMain: Boolean = true): Pair<SourceSet, Co
         configurations[sourceSet.runtimeOnlyConfigurationName]
             .extendsFrom(configurations.runtimeClasspath.get())
     }
+    if (jpsKind != null) {
+        // For Pill
+        sourceSet.extra["jpsKind"] = jpsKind
+    }
     return sourceSet to api
 }
 
@@ -25,7 +29,7 @@ val (builtinsSourceSet, builtinsApi) = extraSourceSet("builtins", extendMain = f
 val (evaluateSourceSet, evaluateApi) = extraSourceSet("evaluate")
 val (interpreterSourceSet, interpreterApi) = extraSourceSet("interpreter")
 val (protobufSourceSet, protobufApi) = extraSourceSet("protobuf")
-val (protobufCompareSourceSet, protobufCompareApi) = extraSourceSet("protobufCompare")
+val (protobufCompareSourceSet, protobufCompareApi) = extraSourceSet("protobufCompare", jpsKind = SourceSet.TEST_SOURCE_SET_NAME)
 val (wasmSourceSet, wasmApi) = extraSourceSet("wasm")
 
 dependencies {
diff --git a/plugins/pill/pill-importer/src/ModelParser.kt b/plugins/pill/pill-importer/src/ModelParser.kt
index 87ae5e35616..60a894d4772 100644
--- a/plugins/pill/pill-importer/src/ModelParser.kt
+++ b/plugins/pill/pill-importer/src/ModelParser.kt
@@ -19,6 +19,7 @@ import org.gradle.api.internal.file.copy.CopySpecInternal
 import org.gradle.api.internal.file.copy.SingleParentCopySpec
 import org.gradle.api.provider.Property
 import org.gradle.jvm.tasks.Jar
+import org.gradle.kotlin.dsl.extra
 import org.gradle.kotlin.dsl.findByType
 import org.gradle.language.jvm.tasks.ProcessResources
 import org.jetbrains.kotlin.pill.model.POrderRoot.*
@@ -349,6 +350,7 @@ private val SourceSet.isTestSourceSet: Boolean
     get() = name == SourceSet.TEST_SOURCE_SET_NAME
             || name.endsWith("Test")
             || name.endsWith("Tests")
+            || (extra.has("jpsKind") && extra.get("jpsKind") == SourceSet.TEST_SOURCE_SET_NAME)
 
 private fun Any.invokeInternal(name: String, instance: Any = this): Any? {
     val method = javaClass.methods.single { it.name.startsWith(name) && it.parameterTypes.isEmpty() }