diff --git a/spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt b/spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt index 375eb16a1e..0965244557 100644 --- a/spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt +++ b/spring-webflux/src/main/kotlin/org/springframework/web/reactive/function/client/WebClientExtensions.kt @@ -16,6 +16,8 @@ package org.springframework.web.reactive.function.client +import kotlinx.coroutines.Job +import kotlinx.coroutines.currentCoroutineContext import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.reactive.asFlow @@ -99,16 +101,20 @@ suspend fun RequestHeadersSpec>.awaitExchange(): Clien * @author Sebastien Deleuze * @since 5.3 */ -suspend fun RequestHeadersSpec>.awaitExchange(responseHandler: suspend (ClientResponse) -> T): T = - exchangeToMono { mono(Dispatchers.Unconfined) { responseHandler.invoke(it) } }.awaitSingle() +suspend fun RequestHeadersSpec>.awaitExchange(responseHandler: suspend (ClientResponse) -> T): T { + val context = currentCoroutineContext().minusKey(Job.Key) + return exchangeToMono { mono(context) { responseHandler.invoke(it) } }.awaitSingle() +} /** * Variant of [WebClient.RequestHeadersSpec.awaitExchange] that allows a nullable return * * @since 5.3.8 */ -suspend fun RequestHeadersSpec>.awaitExchangeOrNull(responseHandler: suspend (ClientResponse) -> T?): T? = - exchangeToMono { mono(Dispatchers.Unconfined) { responseHandler.invoke(it) } }.awaitSingleOrNull() +suspend fun RequestHeadersSpec>.awaitExchangeOrNull(responseHandler: suspend (ClientResponse) -> T?): T? { + val context = currentCoroutineContext().minusKey(Job.Key) + return exchangeToMono { mono(context) { responseHandler.invoke(it) } }.awaitSingleOrNull() +} /** * Coroutines variant of [WebClient.RequestHeadersSpec.exchangeToFlux]. diff --git a/spring-webflux/src/test/kotlin/org/springframework/web/reactive/function/client/WebClientExtensionsTests.kt b/spring-webflux/src/test/kotlin/org/springframework/web/reactive/function/client/WebClientExtensionsTests.kt index 9d6ab8e118..8211dbb7ca 100644 --- a/spring-webflux/src/test/kotlin/org/springframework/web/reactive/function/client/WebClientExtensionsTests.kt +++ b/spring-webflux/src/test/kotlin/org/springframework/web/reactive/function/client/WebClientExtensionsTests.kt @@ -18,7 +18,9 @@ package org.springframework.web.reactive.function.client import io.mockk.every import io.mockk.mockk +import io.mockk.slot import io.mockk.verify +import kotlinx.coroutines.currentCoroutineContext import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.flow import kotlinx.coroutines.flow.toList @@ -32,6 +34,8 @@ import reactor.core.publisher.Flux import reactor.core.publisher.Mono import java.util.concurrent.CompletableFuture import java.util.function.Function +import kotlin.coroutines.AbstractCoroutineContextElement +import kotlin.coroutines.CoroutineContext /** * Mock object based tests for [WebClient] Kotlin extensions @@ -110,6 +114,18 @@ class WebClientExtensionsTests { } } + @Test + fun `awaitExchange with coroutines context`() { + val foo = mockk() + val slot = slot>>() + every { requestBodySpec.exchangeToMono(capture(slot)) } answers { + slot.captured.apply(mockk()) + } + runBlocking(FooContextElement(foo)) { + assertThat(requestBodySpec.awaitExchange { currentCoroutineContext()[FooContextElement]!!.foo }).isEqualTo(foo) + } + } + @Test fun `awaitExchangeOrNull returning null`() { val foo = mockk() @@ -128,6 +144,18 @@ class WebClientExtensionsTests { } } + @Test + fun `awaitExchangeOrNull with coroutines context`() { + val foo = mockk() + val slot = slot>>() + every { requestBodySpec.exchangeToMono(capture(slot)) } answers { + slot.captured.apply(mockk()) + } + runBlocking(FooContextElement(foo)) { + assertThat(requestBodySpec.awaitExchangeOrNull { currentCoroutineContext()[FooContextElement]!!.foo }).isEqualTo(foo) + } + } + @Test fun exchangeToFlow() { val foo = mockk() @@ -209,4 +237,8 @@ class WebClientExtensionsTests { } class Foo + + private data class FooContextElement(val foo: Foo) : AbstractCoroutineContextElement(FooContextElement) { + companion object Key : CoroutineContext.Key + } }