Merge pull request #48 from joeyfeldberg/master

By popular demand I'm merging this change into the plugin :)
fshutdown · Feb 28, 2017 · fb3e353 · fb3e353
2 parents efde4bf + 3649fdc
commit fb3e353
Showing 1 changed file with 26 additions and 11 deletions.
diff --git a/src/main/java/JMeter/plugins/functional/samplers/websocket/WebSocketSampler.java b/src/main/java/JMeter/plugins/functional/samplers/websocket/WebSocketSampler.java
@@ -4,26 +4,30 @@
  */
 package JMeter.plugins.functional.samplers.websocket;
 
-import java.io.IOException;
-
 import org.apache.commons.lang3.StringUtils;
 import org.apache.jmeter.config.Argument;
 import org.apache.jmeter.config.Arguments;
+import org.apache.jmeter.protocol.http.control.Header;
+import org.apache.jmeter.protocol.http.control.HeaderManager;
 import org.apache.jmeter.protocol.http.util.EncoderCache;
 import org.apache.jmeter.protocol.http.util.HTTPArgument;
 import org.apache.jmeter.protocol.http.util.HTTPConstants;
 import org.apache.jmeter.samplers.AbstractSampler;
 import org.apache.jmeter.samplers.Entry;
 import org.apache.jmeter.samplers.SampleResult;
 import org.apache.jmeter.testelement.TestElement;
-import org.apache.jmeter.testelement.property.*;
+import org.apache.jmeter.testelement.TestStateListener;
+import org.apache.jmeter.testelement.property.PropertyIterator;
+import org.apache.jmeter.testelement.property.StringProperty;
+import org.apache.jmeter.testelement.property.TestElementProperty;
 import org.apache.jorphan.logging.LoggingManager;
 import org.apache.jorphan.util.JOrphanUtils;
 import org.apache.log.Logger;
+import org.eclipse.jetty.util.ssl.SslContextFactory;
+import org.eclipse.jetty.websocket.client.ClientUpgradeRequest;
+import org.eclipse.jetty.websocket.client.WebSocketClient;
 
-
-
-
+import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -33,11 +37,6 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
-import org.apache.jmeter.testelement.TestStateListener;
-import org.eclipse.jetty.util.ssl.SslContextFactory;
-import org.eclipse.jetty.websocket.client.ClientUpgradeRequest;
-import org.eclipse.jetty.websocket.client.WebSocketClient;
-
 /**
  *
  * @author Maciej Zaleski
@@ -54,6 +53,8 @@ public class WebSocketSampler extends AbstractSampler implements TestStateListen
     private static final String WS_PREFIX = "ws://"; // $NON-NLS-1$
     private static final String WSS_PREFIX = "wss://"; // $NON-NLS-1$
     private static final String DEFAULT_PROTOCOL = "ws";
+
+    private HeaderManager headerManager;
 
     private static Map<String, ServiceSocket> connectionList;
 
@@ -88,6 +89,13 @@ private ServiceSocket getConnectionSocket() throws URISyntaxException, Exception
         //Start WebSocket client thread and upgrage HTTP connection
         webSocketClient.start();
         ClientUpgradeRequest request = new ClientUpgradeRequest();
+        if (headerManager != null) {
+            for (int i = 0; i < headerManager.size(); i++) {
+                Header header = headerManager.get(i);
+                request.setHeader(header.getName(), header.getValue());
+            }
+        }
+
         webSocketClient.connect(socket, uri, request);
 
         //Get connection timeout or use the default value
@@ -477,6 +485,13 @@ public Arguments getQueryStringParameters() {
         return args;
     }
 
+    public void addTestElement(TestElement el) {
+        if (el instanceof HeaderManager) {
+            headerManager = (HeaderManager) el;
+        } else {
+            super.addTestElement(el);
+        }
+    }
 
     @Override
     public void testStarted() {