From 33fda4e700df9a1bb2ed3b9ff29d92a461ccb663 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@pivotal.io>
Date: Mon, 12 Feb 2018 13:30:50 +0100
Subject: [PATCH] Added an example with custom rest template builder

---
 .../TraceWebClientAutoConfiguration.java      | 21 +++++++++++++------
 .../TraceWebClientAutoConfigurationTests.java | 17 ++++++++++++---
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java
index 5fa16c009..4ea1c6bc1 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java
@@ -81,12 +81,7 @@ public class TraceWebClientAutoConfiguration {
 		@Bean
 		@Order
 		RestTemplateCustomizer traceRestTemplateCustomizer() {
-			return new RestTemplateCustomizer() {
-				@Override public void customize(RestTemplate restTemplate) {
-					new RestTemplateInterceptorInjector(TraceInterceptorConfiguration.this.traceRestTemplateInterceptor)
-							.inject(restTemplate);
-				}
-			};
+			return new TraceRestTemplateCustomizer(this.traceRestTemplateInterceptor);
 		}
 
 		@PostConstruct
@@ -174,6 +169,20 @@ class RestTemplateInterceptorInjector {
 	}
 }
 
+class TraceRestTemplateCustomizer implements RestTemplateCustomizer {
+
+	private final TraceRestTemplateInterceptor interceptor;
+
+	TraceRestTemplateCustomizer(TraceRestTemplateInterceptor interceptor) {
+		this.interceptor = interceptor;
+	}
+
+	@Override public void customize(RestTemplate restTemplate) {
+		new RestTemplateInterceptorInjector(this.interceptor)
+				.inject(restTemplate);
+	}
+}
+
 class TraceUserInfoRestTemplateCustomizer implements UserInfoRestTemplateCustomizer {
 
 	private final BeanFactory beanFactory;
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfigurationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfigurationTests.java
index f49b83711..0fc2add77 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfigurationTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfigurationTests.java
@@ -18,7 +18,6 @@ package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 
 import org.junit.Test;
@@ -49,11 +48,13 @@ public class TraceWebClientAutoConfigurationTests {
 
 	@Autowired @Qualifier("firstRestTemplate") RestTemplate restTemplate;
 	@Autowired @Qualifier("secondRestTemplate") RestTemplate secondRestTemplate;
+	@Autowired RestTemplateBuilder builder;
 
 	@Test
 	public void should_add_rest_template_interceptors() {
-		assertInterceptorsOrder(assertInterceptorsNotEmpty(this.restTemplate));
+ 		assertInterceptorsOrder(assertInterceptorsNotEmpty(this.restTemplate));
 		assertInterceptorsOrder(assertInterceptorsNotEmpty(this.secondRestTemplate));
+		assertInterceptorsOrder(assertInterceptorsNotEmpty(this.builder.build()));
 	}
 
 	private List<ClientHttpRequestInterceptor> assertInterceptorsNotEmpty(RestTemplate restTemplate) {
@@ -88,14 +89,23 @@ public class TraceWebClientAutoConfigurationTests {
 	@EnableAutoConfiguration
 	static class Config {
 
+		// custom builder
+		@Bean
+		RestTemplateBuilder myRestTemplateBuilder(List<RestTemplateCustomizer> customizers) {
+			return new RestTemplateBuilder()
+					.additionalCustomizers(customizers)
+					.additionalInterceptors(new MyClientHttpRequestInterceptor());
+		}
+
+		// rest template from builder
 		@Bean
 		@Qualifier("firstRestTemplate")
 		RestTemplate restTemplate(RestTemplateBuilder restTemplateBuilder) {
 			return restTemplateBuilder
-					.additionalInterceptors(new MyClientHttpRequestInterceptor())
 					.build();
 		}
 
+		// manual rest template
 		@Bean
 		@Qualifier("secondRestTemplate")
 		RestTemplate secondRestTemplate() {
@@ -106,6 +116,7 @@ public class TraceWebClientAutoConfigurationTests {
 			return restTemplate;
 		}
 
+		// custom customizer
 		@Bean
 		RestTemplateCustomizer myRestTemplateCustomizer() {
 			return restTemplate -> {