diff --git a/src/main/java/JMeter/plugins/functional/samplers/websocket/WebSocketSampler.java b/src/main/java/JMeter/plugins/functional/samplers/websocket/WebSocketSampler.java index 2747394..8a2975e 100644 --- a/src/main/java/JMeter/plugins/functional/samplers/websocket/WebSocketSampler.java +++ b/src/main/java/JMeter/plugins/functional/samplers/websocket/WebSocketSampler.java @@ -4,11 +4,11 @@ */ 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; @@ -16,14 +16,18 @@ 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 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() {