Customize RestTemplateBuilder

Author: bedla

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/client/RestTemplateAutoConfiguration.java

@@ -33,12 +33,14 @@
 import org.springframework.boot.autoconfigure.http.HttpMessageConvertersAutoConfiguration;
 import org.springframework.boot.autoconfigure.web.client.RestTemplateAutoConfiguration.NotReactiveWebApplicationCondition;
 import org.springframework.boot.web.client.RestTemplateBuilder;
+import org.springframework.boot.web.client.RestTemplateBuilderCustomizer;
 import org.springframework.boot.web.client.RestTemplateCustomizer;
 import org.springframework.boot.web.client.RestTemplateRequestCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Conditional;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Lazy;
+import org.springframework.util.Assert;
 import org.springframework.web.client.RestTemplate;
 
 /**
@@ -59,6 +61,7 @@ public class RestTemplateAutoConfiguration {
 	@ConditionalOnMissingBean
 	public RestTemplateBuilder restTemplateBuilder(ObjectProvider<HttpMessageConverters> messageConverters,
 			ObjectProvider<RestTemplateCustomizer> restTemplateCustomizers,
+			ObjectProvider<RestTemplateBuilderCustomizer> restTemplateBuilderCustomizer,
 			ObjectProvider<RestTemplateRequestCustomizer<?>> restTemplateRequestCustomizers) {
 		RestTemplateBuilder builder = new RestTemplateBuilder();
 		HttpMessageConverters converters = messageConverters.getIfUnique();
@@ -67,6 +70,10 @@ public RestTemplateBuilder restTemplateBuilder(ObjectProvider<HttpMessageConvert
 		}
 		builder = addCustomizers(builder, restTemplateCustomizers, RestTemplateBuilder::customizers);
 		builder = addCustomizers(builder, restTemplateRequestCustomizers, RestTemplateBuilder::requestCustomizers);
+		for (RestTemplateBuilderCustomizer customizer : restTemplateBuilderCustomizer) {
+			builder = customizer.customize(builder);
+			Assert.notNull(builder, "RestTemplateBuilderCustomizer returned null builder");
+		}
 		return builder;
 	}
 

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/client/RestTemplateAutoConfigurationTests.java

@@ -20,14 +20,17 @@
 import java.util.List;
 
 import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 
+import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.http.HttpMessageConverters;
 import org.springframework.boot.autoconfigure.http.HttpMessageConvertersAutoConfiguration;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.boot.test.context.runner.ReactiveWebApplicationContextRunner;
 import org.springframework.boot.test.context.runner.WebApplicationContextRunner;
 import org.springframework.boot.web.client.RestTemplateBuilder;
+import org.springframework.boot.web.client.RestTemplateBuilderCustomizer;
 import org.springframework.boot.web.client.RestTemplateCustomizer;
 import org.springframework.boot.web.client.RestTemplateRequestCustomizer;
 import org.springframework.context.annotation.Bean;
@@ -42,8 +45,10 @@
 import org.springframework.web.client.RestTemplate;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
@@ -157,6 +162,28 @@ void whenReactiveWebApplicationRestTemplateBuilderIsNotConfigured() {
 				.run((context) -> assertThat(context).doesNotHaveBean(RestTemplateBuilder.class));
 	}
 
+	@Test
+	void customizerShouldCustomizeBuilder() {
+		contextRunner.withUserConfiguration(RestTemplateBuilderCustomizerConfig.class)
+				.run((context) -> {
+					assertThatThrownBy(() -> context.getBean(RestTemplateBuilder.class).build())
+							.isInstanceOf(IllegalStateException.class)
+							.hasMessageContaining("customized builder");
+				});
+	}
+
+	@Test
+	void customizerShouldReturnInstanceNotANull() {
+		contextRunner.withUserConfiguration(RestTemplateBuilderCustomizerReturnsNullConfig.class)
+				.run((context) -> {
+					assertThatThrownBy(() -> context.getBean(RestTemplateBuilder.class).build())
+							.isInstanceOf(BeanCreationException.class)
+							.hasMessageContaining("Error creating bean with name 'restTemplateBuilder'")
+							.hasRootCauseInstanceOf(IllegalArgumentException.class)
+							.hasRootCauseMessage("RestTemplateBuilderCustomizer returned null builder");
+				});
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class RestTemplateConfig {
 
@@ -232,4 +259,23 @@ static class CustomHttpMessageConverter extends StringHttpMessageConverter {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class RestTemplateBuilderCustomizerConfig {
+		@Bean
+		RestTemplateBuilderCustomizer restTemplateBuilderCustomizer() {
+			return oldBuilder -> {
+				final RestTemplateBuilder builder = Mockito.spy(oldBuilder);
+				doThrow(new IllegalStateException("customized builder")).when(builder).build();
+				return builder;
+			};
+		}
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	static class RestTemplateBuilderCustomizerReturnsNullConfig {
+		@Bean
+		RestTemplateBuilderCustomizer restTemplateBuilderCustomizer() {
+			return oldBuilder -> null;
+		}
+	}
 }

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/web/client/RestTemplateBuilderCustomizer.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2012-2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.web.client;
+
+/**
+ * Callback interface that can be used to customize a {@link RestTemplateBuilder}.
+ *
+ * @author Ivo Smid
+ * @see RestTemplateBuilder
+ * @since 5.3.0
+ */
+@FunctionalInterface
+public interface RestTemplateBuilderCustomizer {
+	/**
+	 * Callback to customize a {@link RestTemplateBuilder} instance.
+	 *
+	 * @param restTemplateBuilder the original builder to customize
+	 * @return customized builder instance
+	 */
+	RestTemplateBuilder customize(RestTemplateBuilder restTemplateBuilder);
+}