Support importing main declarations in test source sets

Mr3zee · Mr3zee · commit 84229af38435 · 2025-07-21T20:05:10.000+02:00
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufExecTask.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufExecTask.kt
@@ -14,6 +14,7 @@ import org.gradle.api.provider.ListProperty
 import org.gradle.api.provider.Property
 import org.gradle.api.provider.Provider
 import org.gradle.api.tasks.Input
+import org.gradle.api.tasks.InputDirectory
 import org.gradle.api.tasks.InputFile
 import org.gradle.api.tasks.Optional
 import org.gradle.api.tasks.TaskAction
@@ -36,7 +37,7 @@ public abstract class BufExecTask : DefaultTask() {
     @get:Input
     public abstract val args: ListProperty<String>
 
-    @get:Input
+    @get:InputDirectory
     public abstract val workingDir: Property<File>
 
     @get:InputFile
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufGenYamlUpdate.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufGenYamlUpdate.kt
@@ -5,6 +5,7 @@
 package kotlinx.rpc.buf.tasks
 
 import kotlinx.rpc.buf.BUF_GEN_YAML
+import kotlinx.rpc.proto.PROTO_FILES_DIR
 import kotlinx.rpc.proto.PROTO_GROUP
 import kotlinx.rpc.proto.ProtocPlugin
 import kotlinx.rpc.proto.protoBuildDirSourceSets
@@ -26,7 +27,7 @@ internal data class ResolvedGrpcPlugin(
     val type: Type,
     val locator: List<String>,
     val out: String,
-    val options: Map<String, String?>,
+    val options: Map<String, Any?>,
     val strategy: String?,
     val includeImports: Boolean?,
     val includeWrk: Boolean?,
@@ -59,7 +60,7 @@ public abstract class BufGenYamlUpdate : DefaultTask() {
             file.createNewFile()
         }
 
-        bufGenFile.get().bufferedWriter(Charsets.UTF_8).use { writer ->
+        file.bufferedWriter(Charsets.UTF_8).use { writer ->
             writer.appendLine("version: v2")
             writer.appendLine("clean: true")
             writer.appendLine("plugins:")
@@ -107,6 +108,9 @@ public abstract class BufGenYamlUpdate : DefaultTask() {
                 }
             }
 
+            writer.appendLine("inputs:")
+            writer.appendLine("  - directory: $PROTO_FILES_DIR")
+
             writer.flush()
         }
     }
@@ -155,9 +159,7 @@ internal fun Project.registerBufGenYamlUpdateTask(
         val bufGenYamlFile = project.protoBuildDirSourceSets
             .resolve(dir)
             .resolve(BUF_GEN_YAML)
-            .apply {
-                ensureRegularFileExists()
-            }
+            .ensureRegularFileExists()
 
         bufGenFile.set(bufGenYamlFile)
 
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufGenerateTask.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufGenerateTask.kt
@@ -24,14 +24,13 @@ public abstract class BufGenerateTask : BufExecTask() {
     @get:Optional
     public abstract val additionalArgs: ListProperty<String>
 
+    @get:InputFiles
+    internal abstract val protoFiles: ListProperty<File>
+
     /** Whether to include imports. */
     @get:Input
     internal abstract val includeImports: Property<Boolean>
 
-    /** The input proto files. */
-    @get:InputFiles
-    internal abstract val inputFiles: ConfigurableFileCollection
-
     /** The directory to output generated files. */
     @get:OutputDirectory
     internal abstract val outputDirectory: Property<File>
@@ -57,8 +56,8 @@ public abstract class BufGenerateTask : BufExecTask() {
 internal fun Project.registerBufGenerateTask(
     name: String,
     workingDir: Provider<File>,
-    inputFiles: Provider<FileCollection>,
     outputDirectory: Provider<File>,
+    protoFiles: Provider<FileCollection>,
     configure: BufGenerateTask.() -> Unit = {},
 ): TaskProvider<BufGenerateTask> {
     return registerBufExecTask<BufGenerateTask>(name, workingDir) {
@@ -68,8 +67,8 @@ internal fun Project.registerBufGenerateTask(
         val generate = project.rpcExtension().grpc.buf.generate
 
         includeImports.set(generate.includeImports)
-        this.inputFiles.setFrom(inputFiles)
         this.outputDirectory.set(outputDirectory)
+        this.protoFiles.set(protoFiles)
 
         configure()
     }
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufYamlUpdate.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/buf/tasks/BufYamlUpdate.kt
@@ -5,17 +5,17 @@
 package kotlinx.rpc.buf.tasks
 
 import kotlinx.rpc.buf.BUF_YAML
+import kotlinx.rpc.proto.IMPORT_PROTO_FILES_DIR
 import kotlinx.rpc.proto.PROTO_FILES_DIR
 import kotlinx.rpc.proto.PROTO_GROUP
 import kotlinx.rpc.proto.protoBuildDirSourceSets
 import kotlinx.rpc.util.ensureDirectoryExists
 import kotlinx.rpc.util.ensureRegularFileExists
 import org.gradle.api.DefaultTask
 import org.gradle.api.Project
-import org.gradle.api.provider.ListProperty
 import org.gradle.api.provider.Property
-import org.gradle.api.tasks.Input
 import org.gradle.api.tasks.InputDirectory
+import org.gradle.api.tasks.Optional
 import org.gradle.api.tasks.OutputFile
 import org.gradle.api.tasks.TaskAction
 import org.gradle.api.tasks.TaskProvider
@@ -26,6 +26,10 @@ public abstract class BufYamlUpdate : DefaultTask() {
     @get:InputDirectory
     internal abstract val protoSourceDir: Property<File>
 
+    @get:Optional
+    @get:InputDirectory
+    internal abstract val importSourceDir: Property<File>
+
     @get:OutputFile
     public abstract val bufFile: Property<File>
 
@@ -59,6 +63,12 @@ public abstract class BufYamlUpdate : DefaultTask() {
                 writer.appendLine("  - path: $modulePath")
             }
 
+            val importDir = importSourceDir.orNull
+            if (importDir != null && importDir.exists()) {
+                val modulePath = importDir.relativeTo(file.parentFile)
+                writer.appendLine("  - path: $modulePath")
+            }
+
             writer.flush()
         }
     }
@@ -71,21 +81,29 @@ public abstract class BufYamlUpdate : DefaultTask() {
 internal fun Project.registerBufYamlUpdateTask(
     name: String,
     dir: String,
+    withImport: Boolean,
     configure: BufYamlUpdate.() -> Unit = {},
 ): TaskProvider<BufYamlUpdate> {
     val capitalizeName = name.replaceFirstChar { it.uppercase() }
     return tasks.register<BufYamlUpdate>("${BufYamlUpdate.PREFIX_NAME}$capitalizeName") {
-        val protoDir = project.protoBuildDirSourceSets.resolve(dir).resolve(PROTO_FILES_DIR)
-        protoDir.ensureDirectoryExists()
+        val baseDir = project.protoBuildDirSourceSets.resolve(dir)
+        val protoDir = baseDir
+            .resolve(PROTO_FILES_DIR)
+            .ensureDirectoryExists()
 
         protoSourceDir.set(protoDir)
 
-        val bufYamlFile = project.protoBuildDirSourceSets
-            .resolve(dir)
+        if (withImport) {
+            val importDir = baseDir
+                .resolve(IMPORT_PROTO_FILES_DIR)
+                .ensureDirectoryExists()
+
+            importSourceDir.set(importDir)
+        }
+
+        val bufYamlFile = baseDir
             .resolve(BUF_YAML)
-            .apply {
-                ensureRegularFileExists()
-            }
+            .ensureRegularFileExists()
 
         bufFile.set(bufYamlFile)
 
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/grpc/GrpcExtension.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/grpc/GrpcExtension.kt
@@ -17,7 +17,9 @@ import kotlinx.rpc.buf.tasks.BufGenerateTask
 import kotlinx.rpc.buf.tasks.registerBufGenYamlUpdateTask
 import kotlinx.rpc.buf.tasks.registerBufGenerateTask
 import kotlinx.rpc.buf.tasks.registerBufYamlUpdateTask
+import kotlinx.rpc.proto.IMPORT_PROTO_FILES_DIR
 import kotlinx.rpc.proto.KXRPC_PLUGIN_JAR_CONFIGURATION
+import kotlinx.rpc.proto.PROTO_FILES_DIR
 import kotlinx.rpc.proto.ProtoSourceSet
 import kotlinx.rpc.proto.ProtocPlugin
 import kotlinx.rpc.proto.configureKxRpcPluginJarConfiguration
@@ -89,7 +91,10 @@ public open class GrpcExtension @Inject constructor(
         val baseName = protoSourceSet.baseName.get()
         val baseGenDir = project.protoBuildDirSourceSets.resolve(baseName)
 
-        val protocPluginNames = protoSourceSet.collectProtocPlugins().distinct()
+        val pairSourceSet = protoSourceSet.correspondingMainSourceSetOrNull()
+
+        val mainProtocPlugins = pairSourceSet?.protocPlugins?.get().orEmpty()
+        val protocPluginNames = (protoSourceSet.protocPlugins.get() + mainProtocPlugins).distinct()
 
         val includedProtocPlugins = protocPluginNames.map {
             protocPlugins.findByName(it)
@@ -99,46 +104,81 @@ public open class GrpcExtension @Inject constructor(
         val protoFiles = protoSourceSet.proto
         val hasFiles = !protoFiles.isEmpty
 
-        val bufGenUpdateTask = project.registerBufGenYamlUpdateTask(
+        val bufUpdateTask = registerBufYamlUpdateTask(
+            name = protoSourceSet.name,
+            dir = baseName,
+            withImport = pairSourceSet != null,
+        )
+
+        val bufGenUpdateTask = registerBufGenYamlUpdateTask(
             name = protoSourceSet.name,
             dir = baseName,
             protocPlugins = includedProtocPlugins,
         ) {
-            onlyIf { hasFiles }
-        }
-
-        val bufUpdateTask = project.registerBufYamlUpdateTask(protoSourceSet.name, baseName) {
-            onlyIf { hasFiles }
+            dependsOn(bufUpdateTask)
         }
 
-        val processProtoTask = project.registerProcessProtoFilesTask(
+        val processProtoTask = registerProcessProtoFilesTask(
             name = protoSourceSet.name,
-            baseGenDir = project.provider { baseGenDir },
+            baseGenDir = provider { baseGenDir },
             protoFiles = protoFiles,
+            toDir = PROTO_FILES_DIR,
         ) {
-            onlyIf { hasFiles }
-
-            dependsOn(bufGenUpdateTask)
             dependsOn(bufUpdateTask)
+            dependsOn(bufGenUpdateTask)
+        }
+
+        val processImportProtoTask = if (pairSourceSet != null) {
+            val importProtoFiles = pairSourceSet.proto
+
+            registerProcessProtoFilesTask(
+                name = "${protoSourceSet.name}Import",
+                baseGenDir = provider { baseGenDir },
+                protoFiles = importProtoFiles,
+                toDir = IMPORT_PROTO_FILES_DIR,
+            ) {
+                dependsOn(bufUpdateTask)
+                dependsOn(bufGenUpdateTask)
+                dependsOn(processProtoTask)
+            }
+        } else {
+            null
         }
 
-        val out = project.protoBuildDirGenerated.resolve(baseName)
+        val out = protoBuildDirGenerated.resolve(baseName)
 
         val capitalName = protoSourceSet.name.replaceFirstChar { it.uppercase() }
-        val bufGenerateTask = project.registerBufGenerateTask(
+        val bufGenerateTask = registerBufGenerateTask(
             name = "${BufGenerateTask.NAME_PREFIX}$capitalName",
             workingDir = provider { baseGenDir },
-            inputFiles = processProtoTask.map { it.outputs.files },
             outputDirectory = provider { out },
+            protoFiles = provider {
+                protoFiles.asFileTree.let {
+                    if (pairSourceSet != null) {
+                        it + pairSourceSet.proto
+                    } else {
+                        it
+                    }
+                }
+            },
         ) {
             dependsOn(bufGenUpdateTask)
             dependsOn(bufUpdateTask)
             dependsOn(processProtoTask)
+            if (processImportProtoTask != null) {
+                dependsOn(processImportProtoTask)
+            }
+
+            if (pairSourceSet != null) {
+                dependsOn(pairSourceSet.generateTask)
+            }
 
             onlyIf { hasFiles }
         }
 
-        project.tasks.withType<KotlinCompile>().configureEach {
+        protoSourceSet.generateTask.set(bufGenerateTask)
+
+        tasks.withType<KotlinCompile>().configureEach {
             // compileKotlin - main
             // compileTestKotlin - test
             // compileKotlinJvm - jvmMain
@@ -147,8 +187,8 @@ public open class GrpcExtension @Inject constructor(
             // compileTestKotlinIosArm64 - iosArm64Test
             val taskNameAsSourceSet = name
                 .removePrefix("compile").let {
-                    val suffix = it.substringBefore("Kotlin").takeIf {
-                        prefix -> prefix.isNotEmpty()
+                    val suffix = it.substringBefore("Kotlin").takeIf { prefix ->
+                        prefix.isNotEmpty()
                     } ?: "Main"
 
                     (it.substringAfter("Kotlin")
@@ -162,7 +202,7 @@ public open class GrpcExtension @Inject constructor(
             }
         }
 
-        project.tasks.withType<JavaCompile>().configureEach {
+        tasks.withType<JavaCompile>().configureEach {
             // compileJvmTestJava - test (java, kmp)
             // compileJvmMainJava - main (java, kmp)
             // compileJava - main (java)
@@ -183,7 +223,7 @@ public open class GrpcExtension @Inject constructor(
 
         includedProtocPlugins.forEach { plugin ->
             // locates correctly jvmMain, main jvmTest, test
-            val javaSourceSet = project.extensions.findByType<JavaPluginExtension>()
+            val javaSourceSet = extensions.findByType<JavaPluginExtension>()
                 ?.sourceSets?.findByName(baseName)?.java
 
             if (plugin.isJava.get() && javaSourceSet != null) {
@@ -237,17 +277,17 @@ public open class GrpcExtension @Inject constructor(
         }
     }
 
-    private fun ProtoSourceSet.collectProtocPlugins(): List<String> {
+    private fun ProtoSourceSet.correspondingMainSourceSetOrNull(): ProtoSourceSet? {
         return when {
             name.endsWith("Main") -> {
-                protocPlugins.get()
+                null
             }
 
             name.endsWith("Test") -> {
                 val main = project.protoSourceSets
                     .getByName(correspondingMainName())
 
-                main.collectProtocPlugins()  + protocPlugins.get()
+                main
             }
 
             else -> throw GradleException("Unknown source set name: $name")
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/proto/ProcessProtoFiles.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/proto/ProcessProtoFiles.kt
@@ -22,12 +22,13 @@ internal fun Project.registerProcessProtoFilesTask(
     name: String,
     baseGenDir: Provider<File>,
     protoFiles: SourceDirectorySet,
+    toDir: String,
     configure: ProcessProtoFiles.() -> Unit = {},
 ): TaskProvider<ProcessProtoFiles> {
     val capitalName = name.replaceFirstChar { it.uppercase() }
 
     return tasks.register<ProcessProtoFiles>("process${capitalName}ProtoFiles") {
-        val protoGenDir = baseGenDir.map { it.resolve(PROTO_FILES_DIR) }
+        val protoGenDir = baseGenDir.map { it.resolve(toDir) }
 
         val allFiles = protoFiles.files
 
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/proto/ProtoSourceSet.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/proto/ProtoSourceSet.kt
@@ -4,6 +4,7 @@
 
 package kotlinx.rpc.proto
 
+import kotlinx.rpc.buf.tasks.BufGenerateTask
 import kotlinx.rpc.util.findOrCreate
 import kotlinx.rpc.util.withKotlinJvmExtension
 import kotlinx.rpc.util.withKotlinKmpExtension
@@ -40,6 +41,7 @@ public open class ProtoSourceSet @Inject constructor(internal val project: Proje
     internal val baseName: Property<String> = project.objects.property<String>()
     internal val languageSourceSets: ListProperty<Any> = project.objects.listProperty<Any>()
     internal val protocPlugins: ListProperty<String> = project.objects.listProperty<String>().convention(emptyList())
+    internal val generateTask: Property<BufGenerateTask> = project.objects.property<BufGenerateTask>()
 
     public fun protocPlugin(plugin: NamedDomainObjectProvider<ProtocPlugin>) {
         protocPlugins.add(plugin.name)
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/proto/consts.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/proto/consts.kt
@@ -8,6 +8,7 @@ import org.gradle.api.Project
 
 public const val PROTO_GROUP: String = "proto"
 public const val PROTO_FILES_DIR: String = "proto"
+public const val IMPORT_PROTO_FILES_DIR: String = "import"
 public const val PROTO_BUILD_DIR: String = "protoBuild"
 public const val PROTO_SOURCE_DIRECTORY_NAME: String = "proto"
 public const val PROTO_SOURCE_SETS: String = "protoSourceSets"
diff --git a/gradle-plugin/src/main/kotlin/kotlinx/rpc/util/files.kt b/gradle-plugin/src/main/kotlin/kotlinx/rpc/util/files.kt
diff --git a/grpc/grpc-ktor-server/build.gradle.kts b/grpc/grpc-ktor-server/build.gradle.kts
diff --git a/grpc/grpc-ktor-server/src/jvmMain/proto/some.proto b/grpc/grpc-ktor-server/src/jvmMain/proto/some.proto
diff --git a/grpc/grpc-ktor-server/src/jvmTest/proto/ktor-test-service.proto b/grpc/grpc-ktor-server/src/jvmTest/proto/ktor-test-service.proto
