Unit test & clean up

Author: ssanthanam185

gateway-ha/src/main/java/com/lyft/data/gateway/ha/handler/QueryIdCachingProxyHandler.java

@@ -83,28 +83,11 @@ public void preConnectionHook(HttpServletRequest request, Request proxyRequest)
     }
 
     if (isPathWhiteListed(request.getRequestURI())) {
-      // Set the Host Header to proxy target
       setForwardedHostHeaderOnProxyRequest(request, proxyRequest);
-      logRequestHeaders("After Setting Host Header on proxyRequest", proxyRequest);
     }
 
   }
 
-  private void logRequestHeaders(String msg, Request request) {
-    HttpFields httpFields = request.getHeaders();
-    StringBuilder sb = new StringBuilder();
-    if (httpFields != null) {
-      Enumeration<String> headerNameFields = httpFields.getFieldNames();
-      while (headerNameFields.hasMoreElements()) {
-        String headerName = headerNameFields.nextElement();
-        String value = httpFields.get(headerName);
-        sb.append(headerName).append(":").append(value).append(",");
-      }
-    }
-    log.debug("[{}] For Request [{}], headers are [{}]", msg, request.getURI().toString(),
-        sb.toString());
-  }
-
   private boolean isPathWhiteListed(String path) {
     return path.startsWith(V1_STATEMENT_PATH)
         || path.startsWith(V1_QUERY_PATH)
@@ -282,8 +265,8 @@ protected void postConnectionHook(
     super.postConnectionHook(request, response, buffer, offset, length, callback);
   }
 
-  private void setForwardedHostHeaderOnProxyRequest(HttpServletRequest request,
-                                                    Request proxyRequest) {
+  static void setForwardedHostHeaderOnProxyRequest(HttpServletRequest request,
+                                                   Request proxyRequest) {
     if (request.getHeader(PROXY_TARGET_HEADER) != null) {
       try {
         URI backendUri = new URI(request.getHeader(PROXY_TARGET_HEADER));

gateway-ha/src/test/java/com/lyft/data/gateway/ha/handler/TestQueryIdCachingProxyHandler.java

@@ -4,6 +4,13 @@
 import static org.testng.Assert.assertNull;
 
 import java.io.IOException;
+
+import javax.servlet.http.HttpServletRequest;
+
+import org.eclipse.jetty.client.HttpClient;
+import org.eclipse.jetty.client.api.Request;
+import org.mockito.Mockito;
+import org.testng.Assert;
 import org.testng.annotations.Test;
 
 public class TestQueryIdCachingProxyHandler {
@@ -17,7 +24,7 @@ public void testExtractQueryIdFromUrl() throws IOException {
         "/ui/api/query.html?20200416_160256_03078_6b4yt"};
     for (String path : paths) {
       String queryId = QueryIdCachingProxyHandler.extractQueryIdIfPresent(path, null);
-      assertEquals(queryId,"20200416_160256_03078_6b4yt");
+      assertEquals(queryId, "20200416_160256_03078_6b4yt");
     }
     String[] nonPaths = {
         "/ui/api/query/myOtherThing",
@@ -27,4 +34,20 @@ public void testExtractQueryIdFromUrl() throws IOException {
       assertNull(queryId);
     }
   }
+
+  @Test
+  public void testForwardedHostHeaderOnProxyRequest() throws IOException {
+    String backendServer = "prestocluster";
+    String backendPort = "80";
+    HttpServletRequest mockServletRequest = Mockito.mock(HttpServletRequest.class);
+    Mockito.when(mockServletRequest.getHeader("proxytarget")).thenReturn(String.format("http://%s"
+        + ":%s", backendServer, backendPort));
+    HttpClient httpClient = new HttpClient();
+    Request proxyRequest = httpClient.newRequest("http://localhost:80");
+    QueryIdCachingProxyHandler.setForwardedHostHeaderOnProxyRequest(mockServletRequest,
+        proxyRequest);
+    Assert.assertEquals(proxyRequest.getHeaders().get("Host"), String.format("%s:%s",
+        backendServer, backendPort));
+  }
+
 }