From 80eefb52575d448b430c63edb0a731b0d30c0dd3 Mon Sep 17 00:00:00 2001
From: Joakim Erdfelt <joakim.erdfelt@gmail.com>
Date: Fri, 25 Mar 2016 13:25:08 -0700
Subject: [PATCH] Improving WebSocketAdapter test reliability

---
 .../org/eclipse/jetty/websocket/jsr356/EchoHandler.java     | 1 +
 .../org/eclipse/jetty/websocket/jsr356/JettyEchoSocket.java | 6 ++++++
 .../jetty/websocket/jsr356/MessageReceivingTest.java        | 1 +
 .../org/eclipse/jetty/websocket/api/WebSocketAdapter.java   | 6 ++++--
 4 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/EchoHandler.java b/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/EchoHandler.java
index 9271c288c44..1d82f474d9b 100644
--- a/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/EchoHandler.java
+++ b/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/EchoHandler.java
@@ -31,6 +31,7 @@ public class EchoHandler extends WebSocketHandler implements WebSocketCreator
     @Override
     public void configure(WebSocketServletFactory factory)
     {
+        factory.getPolicy().setMaxTextMessageSize(2 * 1024 * 1024);
         factory.setCreator(this);
     }
 
diff --git a/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/JettyEchoSocket.java b/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/JettyEchoSocket.java
index 3abe1b76355..3e922bcfe2b 100644
--- a/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/JettyEchoSocket.java
+++ b/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/JettyEchoSocket.java
@@ -38,6 +38,9 @@ public class JettyEchoSocket extends WebSocketAdapter
     @Override
     public void onWebSocketBinary(byte[] payload, int offset, int len)
     {
+        if (isNotConnected())
+            return;
+
         try
         {
             RemoteEndpoint remote = getRemote();
@@ -60,6 +63,9 @@ public class JettyEchoSocket extends WebSocketAdapter
     @Override
     public void onWebSocketText(String message)
     {
+        if (isNotConnected())
+            return;
+
         try
         {
             RemoteEndpoint remote = getRemote();
diff --git a/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/MessageReceivingTest.java b/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/MessageReceivingTest.java
index 485ef0a4bf5..935cca187a9 100644
--- a/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/MessageReceivingTest.java
+++ b/jetty-websocket/javax-websocket-client-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/MessageReceivingTest.java
@@ -68,6 +68,7 @@ public class MessageReceivingTest {
     public static void startServer() throws Exception {
         server = new Server();
         ServerConnector connector = new ServerConnector(server);
+        connector.setPort(0);
         server.addConnector(connector);
 
         handler = new EchoHandler();
diff --git a/jetty-websocket/websocket-api/src/main/java/org/eclipse/jetty/websocket/api/WebSocketAdapter.java b/jetty-websocket/websocket-api/src/main/java/org/eclipse/jetty/websocket/api/WebSocketAdapter.java
index 67d413db2eb..73d2c30e297 100644
--- a/jetty-websocket/websocket-api/src/main/java/org/eclipse/jetty/websocket/api/WebSocketAdapter.java
+++ b/jetty-websocket/websocket-api/src/main/java/org/eclipse/jetty/websocket/api/WebSocketAdapter.java
@@ -26,11 +26,11 @@ package org.eclipse.jetty.websocket.api;
 public class WebSocketAdapter implements WebSocketListener
 {
     private volatile Session session;
+    private RemoteEndpoint remote;
 
     public RemoteEndpoint getRemote()
     {
-        Session sess = this.session;
-        return sess == null?null:session.getRemote();
+        return remote;
     }
 
     public Session getSession()
@@ -60,12 +60,14 @@ public class WebSocketAdapter implements WebSocketListener
     public void onWebSocketClose(int statusCode, String reason)
     {
         this.session = null;
+        this.remote = null;
     }
 
     @Override
     public void onWebSocketConnect(Session sess)
     {
         this.session = sess;
+        this.remote = sess.getRemote();
     }
 
     @Override