Add more IT test

Author: arturobernalg

httpclient5-testing/src/test/java/org/apache/hc/client5/testing/websocket/WebSocketClientTest.java

@@ -37,6 +37,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.hc.client5.http.websocket.api.WebSocket;
@@ -51,6 +52,7 @@
 import org.eclipse.jetty.servlet.ServletHolder;
 import org.eclipse.jetty.websocket.api.Session;
 import org.eclipse.jetty.websocket.api.WebSocketAdapter;
+import org.eclipse.jetty.websocket.api.extensions.ExtensionConfig;
 import org.eclipse.jetty.websocket.servlet.WebSocketServlet;
 import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory;
 import org.junit.jupiter.api.AfterEach;
@@ -62,6 +64,7 @@ final class WebSocketClientTest {
 
     private Server server;
     private int port;
+    private final AtomicBoolean pmceNegotiated = new AtomicBoolean(false);
 
     @BeforeEach
     void startServer() throws Exception {
@@ -73,6 +76,7 @@ void startServer() throws Exception {
         final ServletContextHandler ctx = new ServletContextHandler();
         ctx.setContextPath("/");
         ctx.addServlet(new ServletHolder(new EchoServlet()), "/echo");
+        ctx.addServlet(new ServletHolder(new PmceServlet(pmceNegotiated)), "/pmce");
         ctx.addServlet(new ServletHolder(new InterleaveServlet()), "/interleave");
         ctx.addServlet(new ServletHolder(new AbruptServlet()), "/abrupt");
         ctx.addServlet(new ServletHolder(new TooBigServlet()), "/too-big");
@@ -187,6 +191,73 @@ private String buildPayload() {
         }
     }
 
+    @Test
+    void echo_compressed_pmce() throws Exception {
+        final URI uri = uri(port, "/pmce");
+
+        final WebSocketClientConfig cfg = WebSocketClientConfig.custom()
+                .enablePerMessageDeflate(true)
+                .offerServerNoContextTakeover(true)
+                .offerClientNoContextTakeover(true)
+                .offerClientMaxWindowBits(15)
+                .build();
+
+        try (final CloseableWebSocketClient client = WebSocketClients.createDefault()) {
+            client.start();
+
+            final CountDownLatch done = new CountDownLatch(1);
+            final AtomicReference<Throwable> errorRef = new AtomicReference<>();
+            final StringBuilder echoed = new StringBuilder();
+            final AtomicReference<WebSocket> wsRef = new AtomicReference<>();
+
+            final WebSocketListener listener = new WebSocketListener() {
+
+                @Override
+                public void onOpen(final WebSocket ws) {
+                    wsRef.set(ws);
+                    final String payload = "pmce test " + Instant.now();
+                    ws.sendText(payload, true);
+                }
+
+                @Override
+                public void onText(final CharBuffer text, final boolean last) {
+                    echoed.append(text);
+                    if (last) {
+                        final WebSocket ws = wsRef.get();
+                        if (ws != null) {
+                            ws.close(1000, "done");
+                        }
+                    }
+                }
+
+                @Override
+                public void onClose(final int code, final String reason) {
+                    try {
+                        assertEquals(1000, code);
+                        assertTrue(pmceNegotiated.get(), "PMCE not negotiated on server");
+                        assertTrue(echoed.length() > 0, "No text echoed back");
+                    } finally {
+                        done.countDown();
+                    }
+                }
+
+                @Override
+                public void onError(final Throwable ex) {
+                    errorRef.set(ex);
+                    done.countDown();
+                }
+            };
+
+            client.connect(uri, listener, cfg, null);
+            assertTrue(done.await(10, TimeUnit.SECONDS), "WebSocket did not close in time");
+
+            final Throwable error = errorRef.get();
+            if (error != null) {
+                Assertions.fail("WebSocket error: " + error.getMessage(), error);
+            }
+        }
+    }
+
     @Test
     void ping_interleaved_fragmentation() throws Exception {
         final CountDownLatch gotText = new CountDownLatch(1);
@@ -199,12 +270,9 @@ void ping_interleaved_fragmentation() throws Exception {
 
             final URI u = uri(port, "/interleave");
             client.connect(u, new WebSocketListener() {
-                private WebSocket ws;
-
                 @Override
                 public void onOpen(final WebSocket ws) {
                     ws.ping(null);
-                    this.ws = ws;
                     final String prefix = "hello from hc5 WS @ " + Instant.now() + " — ";
                     final StringBuilder sb = new StringBuilder();
                     for (int i = 0; i < 256; i++) {
@@ -339,6 +407,55 @@ public void onWebSocketText(final String msg) {
         }
     }
 
+    public static final class PmceServlet extends WebSocketServlet {
+        private final AtomicBoolean negotiated;
+
+        public PmceServlet(final AtomicBoolean negotiated) {
+            this.negotiated = negotiated;
+        }
+
+        @Override
+        public void configure(final WebSocketServletFactory factory) {
+            factory.getPolicy().setIdleTimeout(30000);
+            factory.setCreator((req, resp) -> new PmceSocket(negotiated));
+        }
+    }
+
+    public static final class PmceSocket extends WebSocketAdapter {
+        private final AtomicBoolean negotiated;
+
+        public PmceSocket(final AtomicBoolean negotiated) {
+            this.negotiated = negotiated;
+        }
+
+        @Override
+        public void onWebSocketConnect(final Session sess) {
+            super.onWebSocketConnect(sess);
+            if (sess != null) {
+                final java.util.List<ExtensionConfig> exts = sess.getUpgradeRequest().getExtensions();
+                boolean hasPmce = false;
+                if (exts != null) {
+                    for (final ExtensionConfig ext : exts) {
+                        if ("permessage-deflate".equalsIgnoreCase(ext.getName())) {
+                            hasPmce = true;
+                            break;
+                        }
+                    }
+                }
+                negotiated.set(hasPmce);
+            }
+        }
+
+        @Override
+        public void onWebSocketText(final String msg) {
+            final Session s = getSession();
+            if (s != null && s.isOpen()) {
+                s.getRemote().sendString(msg, null);
+                s.close(1000, "done");
+            }
+        }
+    }
+
     public static final class InterleaveServlet extends WebSocketServlet {
         @Override
         public void configure(final WebSocketServletFactory factory) {