diff --git a/network-generator/src/main/kotlin/dev/icerock/moko/network/MultiPlatformNetworkGeneratorPlugin.kt b/network-generator/src/main/kotlin/dev/icerock/moko/network/MultiPlatformNetworkGeneratorPlugin.kt
index d7933f0..f41a4bc 100644
--- a/network-generator/src/main/kotlin/dev/icerock/moko/network/MultiPlatformNetworkGeneratorPlugin.kt
+++ b/network-generator/src/main/kotlin/dev/icerock/moko/network/MultiPlatformNetworkGeneratorPlugin.kt
@@ -7,42 +7,39 @@ package dev.icerock.moko.network
 import dev.icerock.moko.network.tasks.GenerateTask
 import org.gradle.api.Plugin
 import org.gradle.api.Project
+import org.gradle.api.Task
+import org.gradle.api.tasks.TaskProvider
+import org.gradle.kotlin.dsl.getByType
 import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
+import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
 import org.jetbrains.kotlin.gradle.tasks.AbstractKotlinCompile
 import org.jetbrains.kotlin.gradle.tasks.AbstractKotlinNativeCompile
 
-@Suppress("ForbiddenComment")
 class MultiPlatformNetworkGeneratorPlugin : Plugin<Project> {
 
-    override fun apply(target: Project) {
-        val mokoNetworkExtension = target.extensions.create("mokoNetwork", SpecConfig::class.java)
+    override fun apply(project: Project) {
+        val mokoNetworkExtension = project.extensions.create("mokoNetwork", SpecConfig::class.java)
 
-        target.afterEvaluate { it.setupProject(mokoNetworkExtension) }
+        project.afterEvaluate {
+            it.setupProject(mokoNetworkExtension)
+        }
     }
 
-    private fun Project.setupProject(extension: SpecConfig) {
-        val multiplatformExtension =
-            extensions.findByType(KotlinMultiplatformExtension::class.java)
+    private fun Project.setupProject(mokoNetworkExtension: SpecConfig) {
+        val multiplatformExtension: KotlinMultiplatformExtension = project.extensions.getByType()
 
-        if (extension.specs.isEmpty()) return
+        if (mokoNetworkExtension.specs.isEmpty()) return
 
         val openApiGenerateTask = tasks.create("openApiGenerate") {
             it.group = "moko-network"
         }
 
-        extension.specs.forEach { spec ->
-            val generatedDir = "$buildDir/generated/moko-network/${spec.name}"
-            val generatedSourcesDir = "$generatedDir/src/main/kotlin"
-            val sourceSet = multiplatformExtension?.sourceSets?.getByName(spec.sourceSet)
-
-            sourceSet?.kotlin?.srcDir(generatedSourcesDir)
-
-            val generateTask: GenerateTask = tasks.create(
-                "${spec.name}OpenApiGenerate",
-                GenerateTask::class.java
-            ) { it.configure(spec, generatedDir) }
-
-            openApiGenerateTask.dependsOn(generateTask)
+        mokoNetworkExtension.specs.forEach { spec ->
+            registerSpecGenerationTask(
+                spec = spec,
+                openApiGenerateTask = openApiGenerateTask,
+                multiplatformExtension = multiplatformExtension
+            )
         }
 
         tasks.matching { it.name == "preBuild" }
@@ -53,3 +50,22 @@ class MultiPlatformNetworkGeneratorPlugin : Plugin<Project> {
             .all { it.dependsOn(openApiGenerateTask) }
     }
 }
+
+private fun Project.registerSpecGenerationTask(
+    spec: SpecInfo,
+    openApiGenerateTask: Task,
+    multiplatformExtension: KotlinMultiplatformExtension,
+) {
+    val generatedDir = "$buildDir/generated/moko-network/${spec.name}"
+    val generatedSourcesDir = "$generatedDir/src/main/kotlin"
+    val sourceSet: KotlinSourceSet? = multiplatformExtension.sourceSets.getByName(spec.sourceSet)
+
+    sourceSet?.kotlin?.srcDir(generatedSourcesDir)
+
+    val generateTask: TaskProvider<GenerateTask> = tasks.register(
+        "${spec.name}OpenApiGenerate",
+        GenerateTask::class.java
+    ) { it.configure(spec, generatedDir) }
+
+    openApiGenerateTask.dependsOn(generateTask)
+}