From e24b2e6b5d8a6028eda923a7aca55a3726b7b4f3 Mon Sep 17 00:00:00 2001
From: Gabriel <gabriel.erzinger@ifood.com.br>
Date: Thu, 8 Apr 2021 14:39:53 -0300
Subject: [PATCH] Add awaitExchangeOrNull extension function to reactive
 webclient

Closes gh-26778
---
 .../function/client/WebClientExtensions.kt     |  9 +++++++++
 .../client/WebClientExtensionsTests.kt         | 18 ++++++++++++++++++
 2 files changed, 27 insertions(+)

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 5b1420182cf..602b71212a7 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
@@ -90,6 +90,15 @@ suspend fun RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchange(): Clien
 suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchange(responseHandler: suspend (ClientResponse) -> T): T =
 		exchangeToMono { mono(Dispatchers.Unconfined) { responseHandler.invoke(it) } }.awaitSingle()
 
+/**
+ * Variant of [WebClient.RequestHeadersSpec.awaitExchange] that allows a nullable return
+ *
+ * @since 5.3.8
+ */
+@Suppress("DEPRECATION")
+suspend fun <T: Any> RequestHeadersSpec<out RequestHeadersSpec<*>>.awaitExchangeOrNull(responseHandler: suspend (ClientResponse) -> T?): T? =
+		exchangeToMono { mono(Dispatchers.Unconfined) { 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 ac127c0709e..7bca596a010 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
@@ -103,6 +103,24 @@ class WebClientExtensionsTests {
 		}
 	}
 
+	@Test
+	fun `awaitExchangeOrNull returning null`() {
+		val foo = mockk<Foo>()
+		every { requestBodySpec.exchangeToMono(any<Function<ClientResponse, Mono<Foo?>>>()) } returns Mono.empty()
+		runBlocking {
+			assertThat(requestBodySpec.awaitExchangeOrNull { foo }).isEqualTo(null)
+		}
+	}
+
+	@Test
+	fun `awaitExchangeOrNull returning object`() {
+		val foo = mockk<Foo>()
+		every { requestBodySpec.exchangeToMono(any<Function<ClientResponse, Mono<Foo>>>()) } returns Mono.just(foo)
+		runBlocking {
+			assertThat(requestBodySpec.awaitExchangeOrNull { foo }).isEqualTo(foo)
+		}
+	}
+
 	@Test
 	fun exchangeToFlow() {
 		val foo = mockk<Foo>()