Fixes #140

Atmosphere · Dec 6, 2016 · 8cbc72f · 8cbc72f
1 parent 1f54e94
commit 8cbc72f
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 21 deletions.
diff --git a/wasync/src/main/java/org/atmosphere/wasync/impl/AtmosphereRequest.java b/wasync/src/main/java/org/atmosphere/wasync/impl/AtmosphereRequest.java
@@ -226,7 +226,11 @@ public AtmosphereRequest build() {
                 _addDecoder(0, bDecoder);
             }
 
-            if (trackMessageLength) {
+            if (trackMessageLength) {                
+                List<String> l = new ArrayList<String>();
+                l.add("true");
+                queryString.put("X-Atmosphere-TrackMessageSize", l);
+
                 TrackMessageSizeDecoder trackMessageSizeDecoder;
                 if (trackMessageLengthDelimiter.length() > 0) {
                     trackMessageSizeDecoder = new TrackMessageSizeDecoder(trackMessageLengthDelimiter, enableProtocol);

diff --git a/wasync/src/test/java/org/atmosphere/tests/BaseTest.java b/wasync/src/test/java/org/atmosphere/tests/BaseTest.java
@@ -15,9 +15,25 @@
  */
 package org.atmosphere.tests;
 
-import com.ning.http.client.AsyncHttpClient;
-import com.ning.http.client.AsyncHttpClientConfig;
-import com.ning.http.client.providers.netty.NettyAsyncHttpProviderConfig;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertTrue;
+
+import java.io.IOException;
+import java.io.Reader;
+import java.io.StringReader;
+import java.net.ConnectException;
+import java.net.ServerSocket;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.atmosphere.client.TrackMessageSizeInterceptor;
 import org.atmosphere.cpr.AtmosphereHandler;
 import org.atmosphere.cpr.AtmosphereResource;
 import org.atmosphere.cpr.AtmosphereResourceEvent;
@@ -41,23 +57,9 @@
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
-import java.io.IOException;
-import java.io.Reader;
-import java.io.StringReader;
-import java.net.ConnectException;
-import java.net.ServerSocket;
-import java.util.Map;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
-
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertNotNull;
-import static org.testng.Assert.assertNotSame;
-import static org.testng.Assert.assertTrue;
+import com.ning.http.client.AsyncHttpClient;
+import com.ning.http.client.AsyncHttpClientConfig;
+import com.ning.http.client.providers.netty.NettyAsyncHttpProviderConfig;
 
 public abstract class BaseTest {
     public final static String RESUME = "Resume";
@@ -312,6 +314,79 @@ public void on(Throwable t) {
 
         assertEquals(response.get(), RESUME);
     }
+
+    @Test
+    public void basicWebSocketWithSizeTrackingOnTest() throws Exception {
+        final CountDownLatch l = new CountDownLatch(1);
+
+        Config config = new Config.Builder()
+                .port(port)
+                .host("127.0.0.1")
+                .resource("/suspend", new AtmosphereHandler() {
+
+                    private final AtomicBoolean b = new AtomicBoolean(false);
+
+                    @Override
+                    public void onRequest(AtmosphereResource r) throws IOException {
+                        if (!b.getAndSet(true)) {
+                            r.suspend(-1);
+                        } else {
+                            r.getBroadcaster().broadcast(RESUME);
+                        }
+                    }
+
+                    @Override
+                    public void onStateChange(AtmosphereResourceEvent r) throws IOException {
+                        if (!r.isResuming() || !r.isCancelled()) {
+                            r.getResource().getResponse().getWriter().print(r.getMessage());
+                            r.getResource().resume();
+                        }
+                    }
+
+                    @Override
+                    public void destroy() {
+
+                    }
+                }).interceptor(new TrackMessageSizeInterceptor()).build();
+
+        server = new Nettosphere.Builder().config(config).build();
+        assertNotNull(server);
+        server.start();
+
+        final CountDownLatch latch = new CountDownLatch(1);
+        final AtomicReference<String> response = new AtomicReference<String>();
+        AtmosphereClient client = ClientFactory.getDefault().newClient(AtmosphereClient.class);
+
+        RequestBuilder request = client.newRequestBuilder()
+                .method(Request.METHOD.GET)
+                .uri(targetUrl + "/suspend").trackMessageLength(true)
+                .transport(transport());
+
+        Socket socket = client.create(client.newOptionsBuilder().runtime(ahc, false).build());
+        ;
+        socket.on("message", new Function<String>() {
+            @Override
+            public void on(String t) {
+                logger.info("Function invoked {}", t);
+                response.set(t);
+                latch.countDown();
+            }
+        }).on(new Function<Throwable>() {
+
+            @Override
+            public void on(Throwable t) {
+                logger.error("", t);
+                latch.countDown();
+            }
+
+        }).open(request.build()).fire("PING");
+
+        latch.await(5, TimeUnit.SECONDS);
+        server.stop();
+        socket.close();
+
+        assertEquals(response.get(), RESUME);
+    }
 
     @Test
     public void allStringFunctionTest() throws Exception {