diff --git a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt index 548203e4776..3acdf1a8850 100644 --- a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt +++ b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt @@ -191,8 +191,10 @@ internal suspend fun Property.awaitFinalValue(): T? { * See also [withRestrictedStages] * * Will ensure that the given [block] can only execute in the given [stage] + * Will wait for the given [stage] if not arrived yet */ internal suspend fun requiredStage(stage: Stage, block: suspend () -> T): T { + if (currentKotlinPluginLifecycle().stage < stage) stage.await() return withRestrictedStages(hashSetOf(stage), block) } diff --git a/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt b/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt index b8227ca0d17..64e2cc87545 100644 --- a/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt +++ b/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt @@ -458,6 +458,16 @@ class KotlinPluginLifecycleTest { } } + @Test + fun `test - requireStage`() = project.runLifecycleAwareTest { + launch { + requiredStage(AfterFinaliseDsl) { + assertEquals(AfterFinaliseDsl, stage) + assertFailsWith { AfterFinaliseDsl.nextOrThrow.await() } + } + } + } + @Test fun `test - withRestrictedStages`() = project.runLifecycleAwareTest { launch {