Merge pull request #164 from william-tran/feature/customize-load-balanced-requests

Customize load balanced requests according to the chosen ServiceInstance

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2015 the original author or authors.
+ * Copyright 2013-2017 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.
@@ -26,7 +26,6 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass;
-import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -39,6 +38,7 @@ import org.springframework.web.client.RestTemplate;
  *
  * @author Spencer Gibb
  * @author Dave Syer
+ * @author Will Tran
  */
 @Configuration
 @ConditionalOnClass(RestTemplate.class)
@@ -65,12 +65,24 @@ public class LoadBalancerAutoConfiguration {
 		};
 	}
 
+	@Autowired(required = false)
+	private List<LoadBalancerRequestTransformer> transformers = Collections.emptyList();
+
+	@Bean
+	@ConditionalOnMissingBean
+	public LoadBalancerRequestFactory loadBalancerRequestFactory(
+			LoadBalancerClient loadBalancerClient) {
+		return new LoadBalancerRequestFactory(loadBalancerClient, transformers);
+	}
+
 	@Configuration
 	@ConditionalOnMissingClass("org.springframework.retry.support.RetryTemplate")
 	static class LoadBalancerInterceptorConfig {
 		@Bean
-		public LoadBalancerInterceptor ribbonInterceptor(LoadBalancerClient loadBalancerClient) {
-			return new LoadBalancerInterceptor(loadBalancerClient);
+		public LoadBalancerInterceptor ribbonInterceptor(
+				LoadBalancerClient loadBalancerClient,
+				LoadBalancerRequestFactory requestFactory) {
+			return new LoadBalancerInterceptor(loadBalancerClient, requestFactory);
 		}
 
 		@Bean
@@ -108,8 +120,10 @@ public class LoadBalancerAutoConfiguration {
 		@Bean
 		public RetryLoadBalancerInterceptor ribbonInterceptor(
 				LoadBalancerClient loadBalancerClient, LoadBalancerRetryProperties properties,
-				LoadBalancedRetryPolicyFactory lbRetryPolicyFactory) {
-			return new RetryLoadBalancerInterceptor(loadBalancerClient, retryTemplate(), properties, lbRetryPolicyFactory);
+				LoadBalancedRetryPolicyFactory lbRetryPolicyFactory,
+				LoadBalancerRequestFactory requestFactory) {
+			return new RetryLoadBalancerInterceptor(loadBalancerClient, retryTemplate(), properties,
+					lbRetryPolicyFactory, requestFactory);
 		}
 
 		@Bean

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerInterceptor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2015 the original author or authors.
+ * Copyright 2013-2017 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.
@@ -18,7 +18,7 @@ package org.springframework.cloud.client.loadbalancer;
 
 import java.io.IOException;
 import java.net.URI;
-import org.springframework.cloud.client.ServiceInstance;
+
 import org.springframework.http.HttpRequest;
 import org.springframework.http.client.ClientHttpRequestExecution;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
@@ -28,13 +28,21 @@ import org.springframework.http.client.ClientHttpResponse;
  * @author Spencer Gibb
  * @author Dave Syer
  * @author Ryan Baxter
+ * @author William Tran
  */
 public class LoadBalancerInterceptor implements ClientHttpRequestInterceptor {
 
 	private LoadBalancerClient loadBalancer;
+	private LoadBalancerRequestFactory requestFactory;
+
+	public LoadBalancerInterceptor(LoadBalancerClient loadBalancer, LoadBalancerRequestFactory requestFactory) {
+		this.loadBalancer = loadBalancer;
+		this.requestFactory = requestFactory;
+	}
 
 	public LoadBalancerInterceptor(LoadBalancerClient loadBalancer) {
-		this.loadBalancer = loadBalancer;
+		// for backwards compatibility
+		this(loadBalancer, new LoadBalancerRequestFactory(loadBalancer));
 	}
 
 	@Override
@@ -42,16 +50,6 @@ public class LoadBalancerInterceptor implements ClientHttpRequestInterceptor {
 			final ClientHttpRequestExecution execution) throws IOException {
 		final URI originalUri = request.getURI();
 		String serviceName = originalUri.getHost();
-		return this.loadBalancer.execute(serviceName,
-				new LoadBalancerRequest<ClientHttpResponse>() {
-					@Override
-					public ClientHttpResponse apply(final ServiceInstance instance)
-							throws Exception {
-						HttpRequest serviceRequest = new ServiceRequestWrapper(request,
-								instance, loadBalancer);
-						return execution.execute(serviceRequest, body);
-					}
-
-				});
+		return this.loadBalancer.execute(serviceName, requestFactory.createRequest(request, body, execution));
 	}
 }

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRequestFactory.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2017 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
+ *
+ *      http://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.cloud.client.loadbalancer;
+
+import java.util.List;
+
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.http.HttpRequest;
+import org.springframework.http.client.ClientHttpRequestExecution;
+import org.springframework.http.client.ClientHttpResponse;
+
+/**
+ * Creates {@link LoadBalancerRequest}s for {@link LoadBalancerInterceptor} and
+ * {@link RetryLoadBalancerInterceptor}. Applies
+ * {@link LoadBalancerRequestTransformer}s to the intercepted
+ * {@link HttpRequest}.
+ * 
+ * @author William Tran
+ *
+ */
+public class LoadBalancerRequestFactory {
+
+	private LoadBalancerClient loadBalancer;
+	private List<LoadBalancerRequestTransformer> transformers;
+
+	public LoadBalancerRequestFactory(LoadBalancerClient loadBalancer,
+			List<LoadBalancerRequestTransformer> transformers) {
+		this.loadBalancer = loadBalancer;
+		this.transformers = transformers;
+	}
+
+	public LoadBalancerRequestFactory(LoadBalancerClient loadBalancer) {
+		this.loadBalancer = loadBalancer;
+	}
+
+	public LoadBalancerRequest<ClientHttpResponse> createRequest(final HttpRequest request,
+			final byte[] body, final ClientHttpRequestExecution execution) {
+		return new LoadBalancerRequest<ClientHttpResponse>() {
+
+			@Override
+			public ClientHttpResponse apply(final ServiceInstance instance)
+					throws Exception {
+				HttpRequest serviceRequest = new ServiceRequestWrapper(request, instance, loadBalancer);
+				if (transformers != null) {
+					for (LoadBalancerRequestTransformer transformer : transformers) {
+						serviceRequest = transformer.transformRequest(serviceRequest, instance);
+					}
+				}
+				return execution.execute(serviceRequest, body);
+			}
+
+		};
+	}
+
+}

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRequestTransformer.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2017 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
+ *
+ *      http://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.cloud.client.loadbalancer;
+
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.core.annotation.Order;
+import org.springframework.http.HttpRequest;
+
+/**
+ * Allows applications to transform the load balanced {@link HttpRequest} given
+ * the chosen {@link ServiceInstance}
+ * 
+ * @author Will Tran
+ */
+@Order(LoadBalancerRequestTransformer.DEFAULT_ORDER)
+public interface LoadBalancerRequestTransformer {
+	public static final int DEFAULT_ORDER = 0;
+
+	HttpRequest transformRequest(HttpRequest request, ServiceInstance instance);
+}

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/RetryLoadBalancerInterceptor.java

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2016-2017 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
+ *
+ *      http://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.cloud.client.loadbalancer;
 
 import java.io.IOException;
@@ -15,6 +31,7 @@ import org.springframework.retry.support.RetryTemplate;
 
 /**
  * @author Ryan Baxter
+ * @author Will Tran
  */
 public class RetryLoadBalancerInterceptor implements ClientHttpRequestInterceptor {
 
@@ -22,15 +39,26 @@ public class RetryLoadBalancerInterceptor implements ClientHttpRequestIntercepto
 	private RetryTemplate retryTemplate;
 	private LoadBalancerClient loadBalancer;
 	private LoadBalancerRetryProperties lbProperties;
+	private LoadBalancerRequestFactory requestFactory;
 
 
 	public RetryLoadBalancerInterceptor(LoadBalancerClient loadBalancer, RetryTemplate retryTemplate,
 										LoadBalancerRetryProperties lbProperties,
-										LoadBalancedRetryPolicyFactory lbRetryPolicyFactory) {
+										LoadBalancedRetryPolicyFactory lbRetryPolicyFactory,
+										LoadBalancerRequestFactory requestFactory) {
 		this.loadBalancer = loadBalancer;
 		this.lbRetryPolicyFactory = lbRetryPolicyFactory;
 		this.retryTemplate = retryTemplate;
 		this.lbProperties = lbProperties;
+		this.requestFactory = requestFactory;
+	}
+	
+	public RetryLoadBalancerInterceptor(LoadBalancerClient loadBalancer, RetryTemplate retryTemplate,
+										LoadBalancerRetryProperties lbProperties,
+										LoadBalancedRetryPolicyFactory lbRetryPolicyFactory) {
+		// for backwards compatibility
+		this(loadBalancer, retryTemplate, lbProperties, lbRetryPolicyFactory, 
+				new LoadBalancerRequestFactory(loadBalancer));
 	}
 
 	@Override
@@ -59,18 +87,7 @@ public class RetryLoadBalancerInterceptor implements ClientHttpRequestIntercepto
 						}
 						return RetryLoadBalancerInterceptor.this.loadBalancer.execute(
 								serviceName, serviceInstance,
-								new LoadBalancerRequest<ClientHttpResponse>() {
-
-									@Override
-									public ClientHttpResponse apply(
-											final ServiceInstance instance)
-											throws Exception {
-										HttpRequest serviceRequest = new ServiceRequestWrapper(
-												request, instance, loadBalancer);
-										return execution.execute(serviceRequest, body);
-									}
-
-								});
+								requestFactory.createRequest(request, body, execution));
 					}
 				});
 	}

spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRequestFactoryConfigurationTests.java

@@ -0,0 +1,172 @@
+/*
+ * Copyright 2017 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
+ *
+ *      http://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.cloud.client.loadbalancer;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.annotation.Order;
+import org.springframework.http.HttpRequest;
+import org.springframework.http.client.ClientHttpRequestExecution;
+
+@RunWith(MockitoJUnitRunner.class)
+public class LoadBalancerRequestFactoryConfigurationTests {
+
+	@Mock
+	private HttpRequest request;
+	@Mock
+	private HttpRequest transformedRequest;
+	@Mock
+	private HttpRequest transformedRequest2;
+	@Mock
+	private ClientHttpRequestExecution execution;
+	@Mock
+	private ServiceInstance instance;
+
+	private byte[] body = new byte[] {};
+	private ArgumentCaptor<HttpRequest> httpRequestCaptor;
+	private LoadBalancerRequestFactory lbReqFactory;
+	private LoadBalancerRequest<?> lbRequest;
+
+	@Before
+	public void setup() {
+		httpRequestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+	}
+
+	protected ConfigurableApplicationContext init(Class<?> config) {
+		ConfigurableApplicationContext context = new SpringApplicationBuilder().web(false)
+				.properties("spring.aop.proxyTargetClass=true")
+				.sources(config, LoadBalancerAutoConfiguration.class).run();
+
+		lbReqFactory = context.getBean(LoadBalancerRequestFactory.class);
+		lbRequest = lbReqFactory.createRequest(request, body, execution);
+		return context;
+	}
+
+	@Test
+	public void transformer() throws Exception {
+		ConfigurableApplicationContext context = init(Transformer.class);
+
+		LoadBalancerRequestTransformer transformer = context.getBean("transformer",
+				LoadBalancerRequestTransformer.class);
+		when(transformer.transformRequest(any(ServiceRequestWrapper.class), eq(instance)))
+				.thenReturn(transformedRequest);
+
+		lbRequest.apply(instance);
+
+		verify(execution).execute(httpRequestCaptor.capture(), eq(body));
+		assertEquals(
+				"transformer should have transformed the ServiceRequestWrapper into transformedRequest",
+				transformedRequest,
+				httpRequestCaptor.getValue());
+	}
+
+	@Test
+	public void noTransformer() throws Exception {
+		init(NoTransformer.class);
+
+		lbRequest.apply(instance);
+
+		verify(execution).execute(httpRequestCaptor.capture(), eq(body));
+		assertEquals(
+				"ServiceRequestWrapper should be executed",
+				ServiceRequestWrapper.class,
+				httpRequestCaptor.getValue().getClass());
+	}
+
+	@Test
+	public void transformersAreOrdered() throws Exception {
+		ConfigurableApplicationContext context = init(TransformersAreOrdered.class);
+
+		LoadBalancerRequestTransformer transformer = context.getBean("transformer",
+				LoadBalancerRequestTransformer.class);
+		when(transformer.transformRequest(any(ServiceRequestWrapper.class), eq(instance)))
+				.thenReturn(transformedRequest);
+		LoadBalancerRequestTransformer transformer2 = context.getBean("transformer2",
+				LoadBalancerRequestTransformer.class);
+		when(transformer2.transformRequest(transformedRequest, instance)).thenReturn(transformedRequest2);
+
+		lbRequest.apply(instance);
+
+		verify(execution).execute(httpRequestCaptor.capture(), eq(body));
+		assertEquals(
+				"transformer2 should run after transformer",
+				transformedRequest2,
+				httpRequestCaptor.getValue());
+	}
+
+	@Configuration
+	static class Transformer {
+
+		@Bean
+		public LoadBalancerClient loadBalancerClient() {
+			return mock(LoadBalancerClient.class);
+		}
+
+		@Bean
+		public LoadBalancerRequestTransformer transformer() {
+			return mock(LoadBalancerRequestTransformer.class);
+		}
+
+	}
+
+	@Configuration
+	static class TransformersAreOrdered {
+
+		@Bean
+		public LoadBalancerClient loadBalancerClient() {
+			return mock(LoadBalancerClient.class);
+		}
+
+		@Bean
+		public LoadBalancerRequestTransformer transformer() {
+			return mock(LoadBalancerRequestTransformer.class);
+		}
+
+		@Bean
+		@Order(LoadBalancerRequestTransformer.DEFAULT_ORDER + 1)
+		public LoadBalancerRequestTransformer transformer2() {
+			return mock(LoadBalancerRequestTransformer.class);
+		}
+
+	}
+
+	@Configuration
+	static class NoTransformer {
+
+		@Bean
+		public LoadBalancerClient loadBalancerClient() {
+			return mock(LoadBalancerClient.class);
+		}
+
+	}
+
+}

spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRequestFactoryTests.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2017 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
+ *
+ *      http://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.cloud.client.loadbalancer;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.http.HttpRequest;
+import org.springframework.http.client.ClientHttpRequestExecution;
+import org.springframework.http.client.ClientHttpResponse;
+
+@RunWith(MockitoJUnitRunner.class)
+public class LoadBalancerRequestFactoryTests {
+
+	@Mock
+	private LoadBalancerClient loadBalancer;
+	@Mock
+	private HttpRequest request;
+	@Mock
+	private HttpRequest transformedRequest1;
+	@Mock
+	private HttpRequest transformedRequest2;
+
+	private byte[] body = new byte[] {};
+
+	@Mock
+	private ClientHttpRequestExecution execution;
+	@Mock
+	private ServiceInstance instance;
+	@Mock
+	private LoadBalancerRequestTransformer transformer1;
+	@Mock
+	private LoadBalancerRequestTransformer transformer2;
+
+	private ArgumentCaptor<HttpRequest> httpRequestCaptor;
+
+	@Before
+	public void setup() {
+		httpRequestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+	}
+
+	@Test
+	public void testNullTransformers() throws Exception {
+		executeLbRequest(null);
+
+		verify(execution).execute(httpRequestCaptor.capture(), eq(body));
+		Assert.assertEquals("request should be of type ServiceRequestWrapper", ServiceRequestWrapper.class,
+				httpRequestCaptor.getValue().getClass());
+	}
+
+	@Test
+	public void testEmptyTransformers() throws Exception {
+		List<LoadBalancerRequestTransformer> transformers = Collections.emptyList();
+
+		executeLbRequest(transformers);
+
+		verify(execution).execute(httpRequestCaptor.capture(), eq(body));
+		Assert.assertEquals("request should be of type ServiceRequestWrapper", ServiceRequestWrapper.class,
+				httpRequestCaptor.getValue().getClass());
+	}
+
+	@Test
+	public void testOneTransformer() throws Exception {
+		List<LoadBalancerRequestTransformer> transformers = Arrays.asList(transformer1);
+		when(transformer1.transformRequest(any(ServiceRequestWrapper.class), eq(instance))).thenReturn(transformedRequest1);
+
+		executeLbRequest(transformers);
+
+		verify(execution).execute(httpRequestCaptor.capture(), eq(body));
+		assertEquals("transformer1 should have transformed request into transformedRequest1", transformedRequest1,
+				httpRequestCaptor.getValue());
+	}
+
+	@Test
+	public void testTwoTransformers() throws Exception {
+		List<LoadBalancerRequestTransformer> transformers = Arrays.asList(transformer1, transformer2);
+		when(transformer1.transformRequest(any(ServiceRequestWrapper.class), eq(instance))).thenReturn(transformedRequest1);
+		when(transformer2.transformRequest(transformedRequest1, instance))
+				.thenReturn(transformedRequest2);
+
+		executeLbRequest(transformers);
+
+		verify(execution).execute(httpRequestCaptor.capture(), eq(body));
+		assertEquals("transformer2 should have transformed transformedRequest1 into transformedRequest2",
+				transformedRequest2,
+				httpRequestCaptor.getValue());
+	}
+
+	private void executeLbRequest(List<LoadBalancerRequestTransformer> transformers) throws Exception {
+		LoadBalancerRequestFactory lbReqFactory = new LoadBalancerRequestFactory(loadBalancer, transformers);
+		LoadBalancerRequest<ClientHttpResponse> lbRequest = lbReqFactory.createRequest(request, body, execution);
+		lbRequest.apply(instance);
+	}
+
+}

spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/RetryLoadBalancerInterceptorTest.java

@@ -35,12 +35,14 @@ public class RetryLoadBalancerInterceptorTest {
     private LoadBalancerClient client;
     private RetryTemplate retryTemplate;
     private LoadBalancerRetryProperties lbProperties;
+    private LoadBalancerRequestFactory lbRequestFactory;
 
     @Before
     public void setUp() throws Exception {
         client = mock(LoadBalancerClient.class);
         retryTemplate = spy(new RetryTemplate());
         lbProperties = new LoadBalancerRetryProperties();
+        lbRequestFactory = mock(LoadBalancerRequestFactory.class);
 
     }
 
@@ -62,11 +64,12 @@ public class RetryLoadBalancerInterceptorTest {
         when(client.choose(eq("foo"))).thenReturn(serviceInstance);
         when(client.execute(eq("foo"), eq(serviceInstance), any(LoadBalancerRequest.class))).thenThrow(new IOException());
         lbProperties.setEnabled(false);
-        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory);
+        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory, lbRequestFactory);
         byte[] body = new byte[]{};
         ClientHttpRequestExecution execution = mock(ClientHttpRequestExecution.class);
         interceptor.intercept(request, body, execution);
         verify(retryTemplate, times(1)).setRetryPolicy(any(NeverRetryPolicy.class));
+        verify(lbRequestFactory).createRequest(request, body, execution);
     }
 
     @Test
@@ -80,11 +83,12 @@ public class RetryLoadBalancerInterceptorTest {
         when(client.choose(eq("foo"))).thenReturn(serviceInstance);
         when(client.execute(eq("foo"), eq(serviceInstance), any(LoadBalancerRequest.class))).thenReturn(clientHttpResponse);
         lbProperties.setEnabled(true);
-        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory);
+        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory, lbRequestFactory);
         byte[] body = new byte[]{};
         ClientHttpRequestExecution execution = mock(ClientHttpRequestExecution.class);
         interceptor.intercept(request, body, execution);
         verify(retryTemplate, times(1)).setRetryPolicy(any(NeverRetryPolicy.class));
+        verify(lbRequestFactory).createRequest(request, body, execution);
     }
 
     @Test
@@ -100,12 +104,13 @@ public class RetryLoadBalancerInterceptorTest {
         when(client.choose(eq("foo"))).thenReturn(serviceInstance);
         when(client.execute(eq("foo"), eq(serviceInstance), any(LoadBalancerRequest.class))).thenReturn(clientHttpResponse);
         lbProperties.setEnabled(true);
-        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory);
+        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory, lbRequestFactory);
         byte[] body = new byte[]{};
         ClientHttpRequestExecution execution = mock(ClientHttpRequestExecution.class);
         ClientHttpResponse rsp = interceptor.intercept(request, body, execution);
         assertThat(rsp, is(clientHttpResponse));
         verify(retryTemplate, times(1)).setRetryPolicy(eq(interceptorRetryPolicy));
+        verify(lbRequestFactory).createRequest(request, body, execution);
     }
 
     @Test
@@ -121,13 +126,14 @@ public class RetryLoadBalancerInterceptorTest {
         when(client.choose(eq("foo"))).thenReturn(serviceInstance);
         when(client.execute(eq("foo"), eq(serviceInstance), any(LoadBalancerRequest.class))).thenThrow(new IOException()).thenReturn(clientHttpResponse);
         lbProperties.setEnabled(true);
-        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory);
+        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory, lbRequestFactory);
         byte[] body = new byte[]{};
         ClientHttpRequestExecution execution = mock(ClientHttpRequestExecution.class);
         ClientHttpResponse rsp = interceptor.intercept(request, body, execution);
         verify(client, times(2)).execute(eq("foo"), eq(serviceInstance), any(LoadBalancerRequest.class));
         assertThat(rsp, is(clientHttpResponse));
         verify(retryTemplate, times(1)).setRetryPolicy(any(InterceptorRetryPolicy.class));
+        verify(lbRequestFactory, times(2)).createRequest(request, body, execution);
     }
 
     @Test(expected = IOException.class)
@@ -144,9 +150,10 @@ public class RetryLoadBalancerInterceptorTest {
         when(client.choose(eq("foo"))).thenReturn(serviceInstance);
         when(client.execute(eq("foo"), eq(serviceInstance), any(LoadBalancerRequest.class))).thenThrow(new IOException()).thenReturn(clientHttpResponse);
         lbProperties.setEnabled(true);
-        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory);
+        RetryLoadBalancerInterceptor interceptor = new RetryLoadBalancerInterceptor(client, retryTemplate, lbProperties, lbRetryPolicyFactory, lbRequestFactory);
         byte[] body = new byte[]{};
         ClientHttpRequestExecution execution = mock(ClientHttpRequestExecution.class);
         ClientHttpResponse rsp = interceptor.intercept(request, body, execution);
+        verify(lbRequestFactory).createRequest(request, body, execution);
     }
 }