YARN-468. coverage fix for org.apache.hadoop.yarn.server.webproxy.amfilter (Aleksey Gorshkov via bobby)

git-svn-id: https://svn.apache.org/repos/asf/hadoop/common/trunk@1456458 13f79535-47bb-0310-9956-ffa450edef68

hadoop-yarn-project/CHANGES.txt

@@ -399,6 +399,8 @@ Release 0.23.7 - UNRELEASED
     YARN-443. allow OS scheduling priority of NM to be different than the 
     containers it launches (tgraves)
 
+    YARN-468. coverage fix for org.apache.hadoop.yarn.server.webproxy.amfilter
+    (Aleksey Gorshkov via bobby)
 
   OPTIMIZATIONS
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-web-proxy/src/test/java/org/apache/hadoop/yarn/server/webproxy/amfilter/TestAmFilter.java

@@ -19,41 +19,39 @@
 package org.apache.hadoop.yarn.server.webproxy.amfilter;
 
 import java.io.IOException;
-import java.util.Collections;
-import java.util.Enumeration;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 
-import javax.servlet.Filter;
-import javax.servlet.FilterChain;
-import javax.servlet.FilterConfig;
-import javax.servlet.ServletContext;
-import javax.servlet.ServletException;
-import javax.servlet.ServletRequest;
-import javax.servlet.ServletResponse;
+import javax.servlet.*;
+import javax.servlet.http.Cookie;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
-import junit.framework.Assert;
+import static junit.framework.Assert.*;
 
+import org.apache.hadoop.yarn.server.webproxy.WebAppProxyServlet;
+import org.glassfish.grizzly.servlet.HttpServletResponseImpl;
 import org.junit.Test;
 import org.mockito.Mockito;
 
+/**
+ * Test AmIpFilter. Requests to a no declared hosts should has way through
+ * proxy. Another requests can be filtered with (without) user name.
+ * 
+ */
+public class TestAmFilter {
 
-public class TestAmFilter  {
-
-  private String proxyHost = "bogushost.com";
+  private String proxyHost = "localhost";
   private String proxyUri = "http://bogus";
+  private String doFilterRequest;
+  private AmIpServletRequestWrapper servletWrapper;
 
   private class TestAmIpFilter extends AmIpFilter {
 
     private Set<String> proxyAddresses = null;
 
     protected Set<String> getProxyAddresses() {
-      if(proxyAddresses == null) {
+      if (proxyAddresses == null) {
         proxyAddresses = new HashSet<String>();
       }
       proxyAddresses.add(proxyHost);
@@ -61,12 +59,10 @@ public class TestAmFilter  {
     }
   }
 
-
   private static class DummyFilterConfig implements FilterConfig {
     final Map<String, String> map;
 
-
-    DummyFilterConfig(Map<String,String> map) {
+    DummyFilterConfig(Map<String, String> map) {
       this.map = map;
     }
 
@@ -74,22 +70,24 @@ public class TestAmFilter  {
     public String getFilterName() {
       return "dummy";
     }
+
     @Override
     public String getInitParameter(String arg0) {
       return map.get(arg0);
     }
+
     @Override
     public Enumeration<String> getInitParameterNames() {
       return Collections.enumeration(map.keySet());
     }
+
     @Override
     public ServletContext getServletContext() {
       return null;
     }
   }
 
-
-  @Test
+  @Test(timeout = 5000)
   public void filterNullCookies() throws Exception {
     HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
 
@@ -97,13 +95,12 @@ public class TestAmFilter  {
     Mockito.when(request.getRemoteAddr()).thenReturn(proxyHost);
 
     HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
-
     final AtomicBoolean invoked = new AtomicBoolean();
 
     FilterChain chain = new FilterChain() {
       @Override
-      public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse)
-        throws IOException, ServletException {
+      public void doFilter(ServletRequest servletRequest,
+          ServletResponse servletResponse) throws IOException, ServletException {
         invoked.set(true);
       }
     };
@@ -115,7 +112,93 @@ public class TestAmFilter  {
     Filter filter = new TestAmIpFilter();
     filter.init(conf);
     filter.doFilter(request, response, chain);
-    Assert.assertTrue(invoked.get());
+    assertTrue(invoked.get());
     filter.destroy();
   }
+
+  /**
+   * Test AmIpFilter
+   */
+  @Test(timeout = 1000)
+  public void testFilter() throws Exception {
+    Map<String, String> params = new HashMap<String, String>();
+    params.put(AmIpFilter.PROXY_HOST, proxyHost);
+    params.put(AmIpFilter.PROXY_URI_BASE, proxyUri);
+    FilterConfig config = new DummyFilterConfig(params);
+
+    // dummy filter
+    FilterChain chain = new FilterChain() {
+      @Override
+      public void doFilter(ServletRequest servletRequest,
+          ServletResponse servletResponse) throws IOException, ServletException {
+        doFilterRequest = servletRequest.getClass().getName();
+        if (servletRequest instanceof AmIpServletRequestWrapper) {
+          servletWrapper = (AmIpServletRequestWrapper) servletRequest;
+
+        }
+      }
+    };
+    AmIpFilter testFilter = new AmIpFilter();
+    testFilter.init(config);
+
+    HttpServletResponseForTest response = new HttpServletResponseForTest();
+    // Test request should implements HttpServletRequest
+
+    ServletRequest failRequest = Mockito.mock(ServletRequest.class);
+    try {
+      testFilter.doFilter(failRequest, response, chain);
+      fail();
+    } catch (ServletException e) {
+      assertEquals("This filter only works for HTTP/HTTPS", e.getMessage());
+    }
+
+    // request with HttpServletRequest
+    HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
+    Mockito.when(request.getRemoteAddr()).thenReturn("redirect");
+    Mockito.when(request.getRequestURI()).thenReturn("/redirect");
+    testFilter.doFilter(request, response, chain);
+    // address "redirect" is not in host list
+    assertEquals("http://bogus/redirect", response.getRedirect());
+    // "127.0.0.1" contains in host list. Without cookie
+    Mockito.when(request.getRemoteAddr()).thenReturn("127.0.0.1");
+    testFilter.doFilter(request, response, chain);
+
+    assertTrue(doFilterRequest
+        .contains("javax.servlet.http.HttpServletRequest"));
+    // cookie added
+    Cookie[] cookies = new Cookie[1];
+    cookies[0] = new Cookie(WebAppProxyServlet.PROXY_USER_COOKIE_NAME, "user");
+
+    Mockito.when(request.getCookies()).thenReturn(cookies);
+    testFilter.doFilter(request, response, chain);
+
+    assertEquals(
+        "org.apache.hadoop.yarn.server.webproxy.amfilter.AmIpServletRequestWrapper",
+        doFilterRequest);
+    // request contains principal from cookie
+    assertEquals("user", servletWrapper.getUserPrincipal().getName());
+    assertEquals("user", servletWrapper.getRemoteUser());
+    assertFalse(servletWrapper.isUserInRole(""));
+
+  }
+
+  private class HttpServletResponseForTest extends HttpServletResponseImpl {
+    String redirectLocation = "";
+
+    public String getRedirect() {
+      return redirectLocation;
+    }
+
+    @Override
+    public void sendRedirect(String location) throws IOException {
+      redirectLocation = location;
+    }
+
+    @Override
+    public String encodeRedirectURL(String url) {
+      return url;
+    }
+
+  }
+
 }