Skip to content

Instantly share code, notes, and snippets.

@d6y

d6y/ws_test.js

Created March 12, 2025 10:46
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save d6y/66d4d32a8b1627d6e17070f7feb5b0b8 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save d6y/66d4d32a8b1627d6e17070f7feb5b0b8 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. d6y revised this gist Mar 12, 2025. No changes.
  2. d6y created this gist Mar 12, 2025.
    96 changes: 96 additions & 0 deletions ws_test.js
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,96 @@
    import http from 'k6/http';
    import ws from 'k6/ws';
    import { check, fail, sleep } from 'k6';

    // App specific settings, whatever your app needs
    let app = {
    apiBase: "http://localhost:8080",
    ws1: "ws://localhost:8080/api/foo/subscribe",
    ws2: "ws://localhost:8080/api/baz/",
    sessionCookieName: "sessionToken",
    userPassword: __ENV.USER_PASSWORD,
    viewSeconds: 10, // how long is a user on the page for?
    };

    export let options = {
    stages: [
    { duration: "1m", target: 60 }, // Stage 1 lasts 1 minute and builds up to 60 users
    { duration: "2m", target: 600 }, // Build up to 600 users
    { duration: "10m", target: 600 }, // Soak test for 10 minutes
    { duration: "2m", target: 0 }, // Ramp down
    ],
    thresholds: {
    http_req_failed: ['rate<0.01'], // http errors should be less than 1%
    http_req_duration: ['p(95)<200'], // 95% of requests should be below 200ms
    },
    };

    // Pre-flight custom validation
    export function setup() {
    if (!app.userPassword) {
    fail("Password for login not set. Run with -e USER_PASSWORD=??? to set the password");
    }
    }


    export default function() {

    // __VU is a "virtual user ID" provided by k6
    const username = `user${__VU}`;

    // This example is grabbing a cookie post-login for authentication:
    const loginResponse = login(username, app.userPassword);
    const cookies = loginResponse.cookies;
    const sessionCookie = cookies[app.sessionCookieName];

    if (!sessionCookie) {
    console.error("No session cookie", username, loginResponse);
    fail("Session cookie not found");
    }

    if (sessionCookie.length !== 1) {
    fail("Expected single cookie value");
    }

    // We will connect to both WS endpoints, then disconnect at the end ("leave the page")

    // We're using a cookie-based authentication scheme
    // The options are limited:
    // https://ably.com/blog/websocket-authentication
    const wsParams = {
    headers: {
    "Cookie": `${sessionCookie[0].name}=${sessionCookie[0].value}`
    }
    };

    // Connect to two web sockets and sit on them for a bit
    const ws1Response = connect(app.ws1, wsParams, function(ws1Socket) {
    const ws2Response = connect(app.ws2, wsParams, function(ws2Socket) {

    sleep(app.viewSeconds);

    ws1Socket ? ws1Socket.close() : fail("ws1 failed");
    ws2Socket ? ws2Socket.close() : fail("ws2 failed");

    });
    check(ws2Response, { '101 ws2 status': (res) => res.status === 101 });
    });

    check(ws1Response, { '101 ws1 status': (res) => res.status === 101 });

    }

    function login(_username, _password) {
    // Exactly how you login is for you to fill in
    // Here we're making a post with some form data and headers
    return http.post(url, formData, { redirects: 0, headers: headers });
    }

    function connect(wsUrl, wsParams, fn) {
    return ws.connect(wsUrl, wsParams, function(socket) {
    socket.on("open", function() {
    fn(socket)
    });
    });
    }