Merge pull request artilleryio#579 from hassy/fix/ws-subprotocol

core/lib/engine_ws.js

@@ -97,28 +97,36 @@ WSEngine.prototype.compile = function (tasks, scenarioSpec, ee) {
 
   return function scenario(initialContext, callback) {
     function zero(callback) {
-      let tls = config.tls || {}; // TODO: config.tls is deprecated
+      let tls = config.tls || {};
       let options = _.extend(tls, config.ws);
 
+      let subprotocols = _.get(config, 'ws.subprotocols', []);
+      const headers = _.get(config, 'ws.headers', {});
+      const subprotocolHeader = _.find(headers, (value, headerName) => {
+        return headerName.toLowerCase() === 'sec-websocket-protocol';
+      });
+      if (typeof subprotocolHeader !== 'undefined') {
+        // NOTE: subprotocols defined via config.ws.subprotocols take precedence:
+        subprotocols = subprotocols.concat(subprotocolHeader.split(',').map(s => s.trim()));
+      }
+
       ee.emit('started');
 
-      let ws = new WebSocket(config.target, options);
+      let ws = new WebSocket(config.target, subprotocols, options);
+
       ws.on('open', function() {
         initialContext.ws = ws;
         return callback(null, initialContext);
       });
+
       ws.once('error', function(err) {
         debug(err);
-        ee.emit('error', err.code);
+        ee.emit('error', err.message || err.code);
         return callback(err, {});
       });
     }
 
     initialContext._successCount = 0;
-    initialContext._pendingRequests = _.size(
-      _.reject(scenarioSpec, function(rs) {
-        return (typeof rs.think === 'number');
-      }));
 
     let steps = _.flatten([
       zero,

test/core/targets/simple_ws.js

@@ -1,5 +1,10 @@
 var WebSocketServer = require('ws').Server;
-var wss = new WebSocketServer({host: '127.0.0.1', port: 9090});
+
+var wss = new WebSocketServer({
+  host: '127.0.0.1',
+  port: parseInt(process.env.PORT) || 9090,
+  handleProtocols: handleProtocols
+});
 
 var MESSAGE_COUNT = 0;
 var CONNECTION_COUNT = 0;
@@ -15,8 +20,13 @@ wss.on('connection', function connection(ws) {
   ws.send('something');
 });
 
-// setInterval(function() {
-//   console.log(new Date());
-//   console.log('CONNECTION_COUNT [ws] = %s', CONNECTION_COUNT);
-//   console.log('MESSAGE_COUNT    [ws] = %s', MESSAGE_COUNT);
-// }, 5 * 1000);
+function handleProtocols(protocols, request) {
+  const SUBPROTOCOL = 'my-custom-protocol';
+  if (protocols.indexOf(SUBPROTOCOL) > -1) {
+    console.log('setting', SUBPROTOCOL);
+    return SUBPROTOCOL;
+  } else {
+    console.log('Unsupported subprotocols', protocols);
+    return false;
+  }
+}

test/core/ws/scripts/subprotocols.json

@@ -0,0 +1,26 @@
+{
+  "config": {
+      "target": "ws://127.0.0.1:9090",
+      "phases": [
+        {"duration": 1, "arrivalCount": 1}
+      ],
+      "variables": {
+        "zipCode": ["35801", "99501", "85001", "94203", "90210"]
+      },
+    "ws": {
+      "subprotocols": ["json", "my-custom-protocol"],
+      "headers": {
+        "sec-websocket-protocol": "my-custom-protocol"
+      }
+    }
+  },
+  "scenarios": [
+    {
+      "engine": "ws",
+      "flow": [
+        {"send": "hello"},
+        {"send": "{{ zipCode }}"}
+      ]
+    }
+  ]
+}

test/core/ws/test_options.js

@@ -30,3 +30,60 @@ test('TLS options for WebSocket', function(t) {
     sessions.run();
   });
 });
+
+test('Subprotocols - using a known subprotocol', function(t) {
+  const script = require('./scripts/subprotocols.json');
+  vuserLauncher(script).then((sessions) => {
+    sessions.on('done', (report) => {
+      t.assert(
+        Object.keys(report.errors).length === 0,
+        'Test with a subprotocol set completed with no errors');
+      t.end();
+    });
+
+    sessions.run();
+  });
+});
+
+test('Subprotocols - no subprotocol', function(t) {
+  const script = require('./scripts/subprotocols.json');
+
+  delete script.config.ws;
+
+  vuserLauncher(script).then((sessions) => {
+    sessions.on('done', (report) => {
+      t.assert(
+        Object.keys(report.errors).length === 0,
+        'Test with no subprotocol set completed with no errors');
+      t.end();
+    });
+
+    sessions.run();
+  });
+});
+
+test('Subprotocols - unknown subprotocol', function(t) {
+  const script = require('./scripts/subprotocols.json');
+
+  script.config.ws = {
+    subprotocols: ['unsupportedByTheServer']
+  };
+
+  vuserLauncher(script).then((sessions) => {
+    sessions.on('done', (report) => {
+      t.assert(
+        Object.keys(report.errors).length === 1,
+        'Should have one error');
+
+      // FIXME: This test is coupled to the error message generated by the ws library
+      t.assert(
+        Object.keys(report.errors)[0].toLowerCase().indexOf('no subprotocol') > -1,
+        'The error should be of "no subprotocol" type'
+      );
+
+      t.end();
+    });
+
+    sessions.run();
+  });
+});