Merge pull request #461 from Ladicek/protocol-upgrade-fix

fix the test for HTTP protocol upgrade

Author: arjantijms

servlet/protocol-handler/pom.xml

@@ -12,4 +12,19 @@
     <packaging>war</packaging>
 
     <name>Java EE 7 Sample: servlet - protocol-handler</name>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <configuration>
+                    <systemPropertyVariables combine.children="append">
+                        <!-- to request protocol upgrade, the client must send the Connection and Upgrade headers -->
+                        <sun.net.http.allowRestrictedHeaders>true</sun.net.http.allowRestrictedHeaders>
+                    </systemPropertyVariables>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
 </project>

servlet/protocol-handler/src/main/java/org/javaee7/servlet/protocolhandler/UpgradeServlet.java

@@ -39,6 +39,7 @@
  */
 package org.javaee7.servlet.protocolhandler;
 
+import static javax.servlet.http.HttpServletResponse.SC_BAD_REQUEST;
 import static javax.servlet.http.HttpServletResponse.SC_SWITCHING_PROTOCOLS;
 
 import java.io.IOException;
@@ -66,11 +67,16 @@ public class UpgradeServlet extends HttpServlet {
      * @throws IOException if an I/O error occurs
      */
     protected void doGet(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
-        response.setStatus(SC_SWITCHING_PROTOCOLS);
-        response.setHeader("Connection", "Upgrade");
-        response.setHeader("Upgrade", "echo");
-        request.upgrade(MyProtocolHandler.class);
-        
-        System.out.println("Request upgraded to MyProtocolHandler");
+        String requestedUpgrade = request.getHeader("Upgrade");
+        if ("echo".equals(requestedUpgrade)) {
+            response.setStatus(SC_SWITCHING_PROTOCOLS);
+            response.setHeader("Connection", "Upgrade");
+            response.setHeader("Upgrade", "echo");
+            request.upgrade(MyProtocolHandler.class);
+
+            System.out.println("Request upgraded to MyProtocolHandler");
+        } else {
+            response.sendError(SC_BAD_REQUEST, "unknown upgrade " + requestedUpgrade);
+        }
     }
 }

servlet/protocol-handler/src/test/java/org/javaee7/servlet/protocolhandler/ProtocolHandlerTest.java

@@ -1,6 +1,6 @@
 package org.javaee7.servlet.protocolhandler;
 
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertEquals;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -44,6 +44,8 @@ public void testUpgradeProtocol() throws IOException, URISyntaxException {
         // typically hang when reading.
 
         URLConnection connection = new URL(base, "UpgradeServlet").openConnection();
+        connection.setRequestProperty("Connection", "Upgrade");
+        connection.setRequestProperty("Upgrade", "echo");
         connection.setConnectTimeout(2000);
         connection.setReadTimeout(2000);
 
@@ -71,7 +73,7 @@ public void testUpgradeProtocol() throws IOException, URISyntaxException {
             }
         }
 
-        assertTrue("In protocol handler".equals(response.toString()));
+        assertEquals("In protocol handler", response.toString());
     }
 
 }