From 671d972454dc594a84070e728f0b8eae1ed2a427 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Deleuze?=
 <sdeleuze@users.noreply.github.com>
Date: Wed, 2 Apr 2025 14:08:08 +0200
Subject: [PATCH] Add RestClient.RequestHeadersSpec#exchangeForRequiredValue

This commit adds a variant to RestClient.RequestHeadersSpec#exchange
suitable for functions returning non-null values.

Closes gh-34692
---
 .../web/client/DefaultRestClient.java         |  7 ++
 .../web/client/RestClient.java                | 75 +++++++++++++++++++
 .../client/RestClientIntegrationTests.java    | 36 ++++++++-
 3 files changed, 117 insertions(+), 1 deletion(-)

diff --git a/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java b/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java
index b877ec388e1..d022774e700 100644
--- a/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java
+++ b/spring-web/src/main/java/org/springframework/web/client/DefaultRestClient.java
@@ -533,6 +533,13 @@ final class DefaultRestClient implements RestClient {
 			return exchangeInternal(exchangeFunction, close);
 		}
 
+		@Override
+		public <T> T exchangeForRequiredValue(RequiredValueExchangeFunction<T> exchangeFunction, boolean close) {
+			T value = exchangeInternal(exchangeFunction, close);
+			Assert.state(value != null, "The exchanged value must not be null");
+			return value;
+		}
+
 		@Nullable
 		private <T> T exchangeInternal(ExchangeFunction<T> exchangeFunction, boolean close) {
 			Assert.notNull(exchangeFunction, "ExchangeFunction must not be null");
diff --git a/spring-web/src/main/java/org/springframework/web/client/RestClient.java b/spring-web/src/main/java/org/springframework/web/client/RestClient.java
index 475f8c864b4..af9c6722e5e 100644
--- a/spring-web/src/main/java/org/springframework/web/client/RestClient.java
+++ b/spring-web/src/main/java/org/springframework/web/client/RestClient.java
@@ -671,12 +671,41 @@ public interface RestClient {
 		 * @param exchangeFunction the function to handle the response with
 		 * @param <T> the type the response will be transformed to
 		 * @return the value returned from the exchange function, potentially {@code null}
+		 * @see RequestHeadersSpec#exchangeForRequiredValue(RequiredValueExchangeFunction)
 		 */
 		@Nullable
 		default <T> T exchange(ExchangeFunction<T> exchangeFunction) {
 			return exchange(exchangeFunction, true);
 		}
 
+		/**
+		 * Exchange the {@link ClientHttpResponse} for a value of type {@code T}.
+		 * This can be useful for advanced scenarios, for example to decode the
+		 * response differently depending on the response status:
+		 * <pre class="code">
+		 * Person person = client.get()
+		 *     .uri("/people/1")
+		 *     .accept(MediaType.APPLICATION_JSON)
+		 *     .exchange((request, response) -&gt; {
+		 *         if (response.getStatusCode().equals(HttpStatus.OK)) {
+		 *             return deserialize(response.getBody());
+		 *         }
+		 *         else {
+		 *             throw new BusinessException();
+		 *         }
+		 *     });
+		 * </pre>
+		 * <p><strong>Note:</strong> The response is
+		 * {@linkplain ClientHttpResponse#close() closed} after the exchange
+		 * function has been invoked.
+		 * @param exchangeFunction the function to handle the response with
+		 * @param <T> the type the response will be transformed to
+		 * @return the value returned from the exchange function, never {@code null}
+		 */
+		default <T> T exchangeForRequiredValue(RequiredValueExchangeFunction<T> exchangeFunction) {
+			return exchangeForRequiredValue(exchangeFunction, true);
+		}
+
 		/**
 		 * Exchange the {@link ClientHttpResponse} for a value of type {@code T}.
 		 * This can be useful for advanced scenarios, for example to decode the
@@ -703,10 +732,40 @@ public interface RestClient {
 		 * {@code exchangeFunction} is invoked, {@code false} to keep it open
 		 * @param <T> the type the response will be transformed to
 		 * @return the value returned from the exchange function, potentially {@code null}
+		 * @see RequestHeadersSpec#exchangeForRequiredValue(RequiredValueExchangeFunction, boolean)
 		 */
 		@Nullable
 		<T> T exchange(ExchangeFunction<T> exchangeFunction, boolean close);
 
+		/**
+		 * Exchange the {@link ClientHttpResponse} for a value of type {@code T}.
+		 * This can be useful for advanced scenarios, for example to decode the
+		 * response differently depending on the response status:
+		 * <pre class="code">
+		 * Person person = client.get()
+		 *     .uri("/people/1")
+		 *     .accept(MediaType.APPLICATION_JSON)
+		 *     .exchange((request, response) -&gt; {
+		 *         if (response.getStatusCode().equals(HttpStatus.OK)) {
+		 *             return deserialize(response.getBody());
+		 *         }
+		 *         else {
+		 *             throw new BusinessException();
+		 *         }
+		 *     });
+		 * </pre>
+		 * <p><strong>Note:</strong> If {@code close} is {@code true},
+		 * then the response is {@linkplain ClientHttpResponse#close() closed}
+		 * after the exchange function has been invoked. When set to
+		 * {@code false}, the caller is responsible for closing the response.
+		 * @param exchangeFunction the function to handle the response with
+		 * @param close {@code true} to close the response after
+		 * {@code exchangeFunction} is invoked, {@code false} to keep it open
+		 * @param <T> the type the response will be transformed to
+		 * @return the value returned from the exchange function, never {@code null}
+		 */
+		<T> T exchangeForRequiredValue(RequiredValueExchangeFunction<T> exchangeFunction, boolean close);
+
 
 		/**
 		 * Defines the contract for {@link #exchange(ExchangeFunction)}.
@@ -726,6 +785,22 @@ public interface RestClient {
 			T exchange(HttpRequest clientRequest, ConvertibleClientHttpResponse clientResponse) throws IOException;
 		}
 
+		/**
+		 * Variant of {@link ExchangeFunction} returning a non-null required value.
+		 * @param <T> the type the response will be transformed to
+		 */
+		@FunctionalInterface
+		interface RequiredValueExchangeFunction<T> extends ExchangeFunction<T> {
+
+			/**
+			 * Exchange the given response into a value of type {@code T}.
+			 * @param clientRequest the request
+			 * @param clientResponse the response
+			 * @return the exchanged value, never {@code null}
+			 * @throws IOException in case of I/O errors
+			 */
+			T exchange(HttpRequest clientRequest, ConvertibleClientHttpResponse clientResponse) throws IOException;
+		}
 
 		/**
 		 * Extension of {@link ClientHttpResponse} that can convert the body.
diff --git a/spring-web/src/test/java/org/springframework/web/client/RestClientIntegrationTests.java b/spring-web/src/test/java/org/springframework/web/client/RestClientIntegrationTests.java
index 916116fa45d..d33dd2d32e2 100644
--- a/spring-web/src/test/java/org/springframework/web/client/RestClientIntegrationTests.java
+++ b/spring-web/src/test/java/org/springframework/web/client/RestClientIntegrationTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -58,6 +58,7 @@ import org.springframework.web.testfixture.xml.Pojo;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 import static org.junit.jupiter.api.Assumptions.assumeFalse;
 import static org.junit.jupiter.params.provider.Arguments.argumentSet;
 
@@ -766,6 +767,39 @@ class RestClientIntegrationTests {
 		expectRequest(request -> assertThat(request.getPath()).isEqualTo("/greeting"));
 	}
 
+	@ParameterizedRestClientTest
+	void exchangeForRequiredValue(ClientHttpRequestFactory requestFactory) {
+		startServer(requestFactory);
+
+		prepareResponse(response -> response.setBody("Hello Spring!"));
+
+		String result = this.restClient.get()
+				.uri("/greeting")
+				.header("X-Test-Header", "testvalue")
+				.exchangeForRequiredValue((request, response) -> new String(RestClientUtils.getBody(response), UTF_8));
+
+		assertThat(result).isEqualTo("Hello Spring!");
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getHeader("X-Test-Header")).isEqualTo("testvalue");
+			assertThat(request.getPath()).isEqualTo("/greeting");
+		});
+	}
+
+	@ParameterizedRestClientTest
+	@SuppressWarnings("DataFlowIssue")
+	void exchangeForNullRequiredValue(ClientHttpRequestFactory requestFactory) {
+		startServer(requestFactory);
+
+		prepareResponse(response -> response.setBody("Hello Spring!"));
+
+		assertThatIllegalStateException().isThrownBy(() -> this.restClient.get()
+				.uri("/greeting")
+				.header("X-Test-Header", "testvalue")
+				.exchangeForRequiredValue((request, response) -> null));
+	}
+
 	@ParameterizedRestClientTest
 	void requestInitializer(ClientHttpRequestFactory requestFactory) {
 		startServer(requestFactory);