diff --git a/ribbon-httpclient/src/test/java/com/netflix/niws/client/http/RestClientTest.java b/ribbon-httpclient/src/test/java/com/netflix/niws/client/http/RestClientTest.java index bcccb6c5..75e989b5 100644 --- a/ribbon-httpclient/src/test/java/com/netflix/niws/client/http/RestClientTest.java +++ b/ribbon-httpclient/src/test/java/com/netflix/niws/client/http/RestClientTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; @@ -137,6 +138,33 @@ public void testDelete() throws Exception { assertStatusIsOk(response.getStatus()); } + @Test + public void testExecuteWithSetLBKey() throws Exception { + final String srcKey = "66666"; + RestClient client = (RestClient) ClientFactory.getNamedClient("allservices"); + BaseLoadBalancer lb = new BaseLoadBalancer() { + public Server chooseServer(Object key) { + if (key != null) { + assertEquals(srcKey, key.toString()); + } + return super.chooseServer(key); + } + }; + final Server[] servers = new Server[]{new Server("localhost", server.getServerPort())}; + lb.addServers(Arrays.asList(servers)); + client.setLoadBalancer(lb); + assertEquals(200, client.executeWithLoadBalancer(HttpRequest.newBuilder().uri(new URI("/")).build()).getStatus()); + assertEquals(200, client.executeWithLoadBalancer(HttpRequest.newBuilder().uri(new URI("/")).loadBalancerKey(srcKey).build()).getStatus()); + client.setLoadBalancer(new BaseLoadBalancer() { + public Server chooseServer(Object key) { + assertNotEquals(srcKey, key.toString()); + addServers(Arrays.asList(servers)); + return super.chooseServer(key); + } + }); + assertEquals(200, client.executeWithLoadBalancer(HttpRequest.newBuilder().uri(new URI("/")).loadBalancerKey("77777").build()).getStatus()); + } + private void assertStatusIsOk(int status) { assertTrue(status == 200 || status == 302); } diff --git a/ribbon-loadbalancer/src/main/java/com/netflix/client/AbstractLoadBalancerAwareClient.java b/ribbon-loadbalancer/src/main/java/com/netflix/client/AbstractLoadBalancerAwareClient.java index 9c9b21d4..0eb96d63 100644 --- a/ribbon-loadbalancer/src/main/java/com/netflix/client/AbstractLoadBalancerAwareClient.java +++ b/ribbon-loadbalancer/src/main/java/com/netflix/client/AbstractLoadBalancerAwareClient.java @@ -128,7 +128,8 @@ protected LoadBalancerCommand buildLoadBalancerCommand(final S request, final LoadBalancerCommand.Builder builder = LoadBalancerCommand.builder() .withLoadBalancerContext(this) .withRetryHandler(handler) - .withLoadBalancerURI(request.getUri()); + .withLoadBalancerURI(request.getUri()) + .withServerLocator(request.getLoadBalancerKey()); customizeLoadBalancerCommandBuilder(request, config, builder); return builder.build(); }