customRoutePlanner: add implementation

Maryum Styles · Maryum Styles · commit f8af0d018f22 · 2021-04-27T17:45:11.000-07:00
implenting customRoutePlanner with the code
from I89591e7c3452be2dc65df9fb705eb54356f61e3c

Bug: T278385
Change-Id: I17a594d7fc36d185add884075d492c07b85b4f6b
diff --git a/http-client-utils/src/main/java/org/wikimedia/utils/http/CustomRoutePlanner.java b/http-client-utils/src/main/java/org/wikimedia/utils/http/CustomRoutePlanner.java
@@ -1,5 +1,56 @@
 package org.wikimedia.utils.http;
 
-public class CustomRoutePlanner {
-    // TODO: stub
+import static java.util.Arrays.stream;
+import static java.util.stream.Collectors.toMap;
+
+import java.util.Map;
+
+import org.apache.http.HttpException;
+import org.apache.http.HttpHost;
+import org.apache.http.HttpRequest;
+import org.apache.http.conn.routing.HttpRoute;
+import org.apache.http.conn.routing.HttpRoutePlanner;
+import org.apache.http.protocol.HttpContext;
+
+public class CustomRoutePlanner implements HttpRoutePlanner {
+
+    private final Map<String, HttpHost> customRouteMap;
+    private final HttpRoutePlanner defaultRoutePlanner;
+
+    public CustomRoutePlanner(Map<String, HttpHost> customRouteMap, HttpRoutePlanner defaultRoutePlanner) {
+        this.customRouteMap = customRouteMap;
+        this.defaultRoutePlanner = defaultRoutePlanner;
+    }
+
+    /**
+     * @param mapProxyProperty has a format of url=url
+     * separated by commas with no space between them
+     * for example: www.wikidata.org=http://proxy.local:9999,www.metawiki.org=https://proxy.local:8080
+     */
+    public static Map<String, HttpHost> createMapFromString(String mapProxyProperty) {
+        String[] pairs = mapProxyProperty.split(",");
+        return stream(pairs).map(p -> p.split("=")).
+                collect(toMap(pair -> pair[0], pair -> HttpHost.create(pair[1])));
+    }
+
+    public void addRoute(String sourceURL, String targetURL) {
+        customRouteMap.put(sourceURL, HttpHost.create(targetURL));
+    }
+
+    @Override
+    public HttpRoute determineRoute(HttpHost httpHost, HttpRequest httpRequest, HttpContext httpContext) throws HttpException {
+        HttpHost destHost = customRouteMap.get(httpHost.getHostName());
+        if (destHost != null) {
+            return new HttpRoute(
+                    new HttpHost(
+                            destHost.getHostName(),
+                            // If the dest port was not set, then assume we want use the same one as the request url.
+                            destHost.getPort() != -1 ? destHost.getPort() : httpHost.getPort(),
+                            destHost.getSchemeName()
+                    )
+            );
+        } else {
+            return defaultRoutePlanner.determineRoute(httpHost, httpRequest, httpContext);
+        }
+    }
 }
diff --git a/http-client-utils/src/test/java/org/wikimedia/utils/http/CustomRoutePlannerUnitTest.java b/http-client-utils/src/test/java/org/wikimedia/utils/http/CustomRoutePlannerUnitTest.java
@@ -2,11 +2,65 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.http.HttpException;
+import org.apache.http.HttpHost;
+import org.apache.http.client.methods.HttpGet;
+import org.apache.http.conn.routing.HttpRoute;
+import org.apache.http.conn.routing.HttpRoutePlanner;
+import org.apache.http.impl.conn.DefaultRoutePlanner;
+import org.apache.http.impl.conn.DefaultSchemePortResolver;
+import org.apache.http.protocol.BasicHttpContext;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 class CustomRoutePlannerUnitTest {
+
+    Map<String, HttpHost> customRouteMapTest = new HashMap<>();
+    HttpRoutePlanner defaultRoutePlanner = new DefaultRoutePlanner(DefaultSchemePortResolver.INSTANCE);
+
+    @BeforeEach
+    public void setUp() {
+        customRouteMapTest.put("managed.test", new HttpHost("route.test", 9999, "someScheme"));
+        customRouteMapTest.put("proxy.test", new HttpHost("route.test", 8080, "diffScheme"));
+    }
+
+    @Test
+    void shouldCreateMapFromString() {
+        String mapString = "www.wikidata.org=http://proxy.local:9999,www.metawiki.org=https://proxy.local:8080";
+        Map<String, HttpHost> stringMap = new HashMap<>();
+        stringMap.put("www.wikidata.org", new HttpHost("proxy.local", 9999, "http"));
+        stringMap.put("www.metawiki.org", new HttpHost("proxy.local", 8080, "https"));
+        assertThat(CustomRoutePlanner.createMapFromString(mapString)).isEqualTo(stringMap);
+    }
+
+
+    @Test
+    void shouldNotGetSameHost() throws HttpException {
+        Map<String, HttpHost> emptyMap = new HashMap<>();
+        CustomRoutePlanner customRoutePlanner = new CustomRoutePlanner(emptyMap, defaultRoutePlanner);
+        HttpHost testHost = new HttpHost("wikidata.org", 9999, "someScheme");
+        HttpRoute actualRoute = customRoutePlanner.determineRoute(testHost, new HttpGet(), new BasicHttpContext());
+        assertThat(actualRoute.getTargetHost()).isEqualTo(testHost);
+    }
+
+    @Test
+    void shouldGetCustomRoute() throws HttpException {
+        CustomRoutePlanner customRoutePlanner = new CustomRoutePlanner(customRouteMapTest, defaultRoutePlanner);
+        HttpHost testHost = new HttpHost("managed.test", 9999, "someScheme");
+        HttpHost expectedHost = new HttpHost("route.test", 9999, "someScheme");
+        HttpRoute actualRoute = customRoutePlanner.determineRoute(testHost, new HttpGet(), new BasicHttpContext());
+        assertThat(actualRoute.getTargetHost()).isEqualTo(expectedHost);
+    }
+
     @Test
-    void test() {
-        assertThat(new CustomRoutePlanner()).isInstanceOf(CustomRoutePlanner.class);
+    void shouldUseProxyRoute() throws HttpException {
+        CustomRoutePlanner customRoutePlanner = new CustomRoutePlanner(customRouteMapTest, defaultRoutePlanner);
+        HttpHost testHost = new HttpHost("proxy.test", 9999, "someScheme");
+        HttpHost expectedHost = new HttpHost("route.test", 8080, "diffScheme");
+        HttpRoute actualRoute = customRoutePlanner.determineRoute(testHost, new HttpGet(), new BasicHttpContext());
+        assertThat(actualRoute.getTargetHost()).isEqualTo(expectedHost);
     }
 }
