Merge branch '6.1.x'

Author: sdeleuze

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<out RequestHeadersSpec<*>>.awaitExchange(): Clien
  * @author Sebastien Deleuze
  * @since 5.3
  */
-suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchange(responseHandler: suspend (ClientResponse) -> T): T =
-		exchangeToMono { mono(Dispatchers.Unconfined) { responseHandler.invoke(it) } }.awaitSingle()
+suspend fun <T: Any> RequestHeadersSpec<out 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 <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchangeOrNull(responseHandler: suspend (ClientResponse) -> T?): T? =
-		exchangeToMono { mono(Dispatchers.Unconfined) { responseHandler.invoke(it) } }.awaitSingleOrNull()
+suspend fun <T: Any> RequestHeadersSpec<out 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].

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<Foo>()
+		val slot = slot<Function<ClientResponse, Mono<Foo>>>()
+		every { requestBodySpec.exchangeToMono(capture(slot)) } answers {
+			slot.captured.apply(mockk<ClientResponse>())
+		}
+		runBlocking(FooContextElement(foo)) {
+			assertThat(requestBodySpec.awaitExchange { currentCoroutineContext()[FooContextElement]!!.foo }).isEqualTo(foo)
+		}
+	}
+
 	@Test
 	fun `awaitExchangeOrNull returning null`() {
 		val foo = mockk<Foo>()
@@ -128,6 +144,18 @@ class WebClientExtensionsTests {
 		}
 	}
 
+	@Test
+	fun `awaitExchangeOrNull with coroutines context`() {
+		val foo = mockk<Foo>()
+		val slot = slot<Function<ClientResponse, Mono<Foo>>>()
+		every { requestBodySpec.exchangeToMono(capture(slot)) } answers {
+			slot.captured.apply(mockk<ClientResponse>())
+		}
+		runBlocking(FooContextElement(foo)) {
+			assertThat(requestBodySpec.awaitExchangeOrNull { currentCoroutineContext()[FooContextElement]!!.foo }).isEqualTo(foo)
+		}
+	}
+
 	@Test
 	fun exchangeToFlow() {
 		val foo = mockk<Foo>()
@@ -209,4 +237,8 @@ class WebClientExtensionsTests {
 	}
 
 	class Foo
+
+	private data class FooContextElement(val foo: Foo) : AbstractCoroutineContextElement(FooContextElement) {
+		companion object Key : CoroutineContext.Key<FooContextElement>
+	}
 }