Automating WebSocket tests

test("echo greets with a welcome message", async () => {
  const s = connect("/ws/echo");
  await s.opened;
  const welcome = JSON.parse(await s.next());

  assert.equal(welcome.type, "welcome");
  assert.match(welcome.message, /echo/i);
  s.ws.close();
});

test("echo returns text, JSON and binary unchanged", async () => {
  const s = connect("/ws/echo");
  await s.opened;
  await s.next(); // welcome

  s.ws.send("héllo 👋");
  assert.equal(await s.next(), "héllo 👋");

  const json = JSON.stringify({ id: 1, tags: ["a", "b"] });
  s.ws.send(json);
  assert.equal(await s.next(), json);

  s.ws.send(new Uint8Array([0, 1, 254, 255]));
  const binary = await s.next();
  assert.ok(binary instanceof ArrayBuffer, "binary should come back as binary");
  assert.deepEqual([...new Uint8Array(binary)], [0, 1, 254, 255]);
  s.ws.close();
});

test("sending close ends the connection with 1000 Bye!", async () => {
  const s = connect("/ws/echo");
  await s.opened;
  await s.next();
  s.ws.send("close");

  const e = await s.closedWithin(3000);
  assert.equal(e.code, 1000);
  assert.equal(e.reason, "Bye!");
  assert.equal(e.wasClean, true);
});

test("a message one byte over 64KB is refused with 1009", async () => {
  const s = connect("/ws/echo");
  await s.opened;
  await s.next();
  s.ws.send("a".repeat(65537));

  const e = await s.closedWithin(3000);
  assert.equal(e.code, 1009);
});

test("ticker sends count ticks in order, then completes and closes", async () => {
  const s = connect("/ws/ticker?count=3");
  await s.opened;

  const ticks = [];
  for (let i = 0; i < 3; i++) ticks.push(JSON.parse(await s.next()));
  assert.deepEqual(ticks.map((t) => t.type), ["tick", "tick", "tick"]);
  assert.deepEqual(ticks.map((t) => t.data.tick), [1, 2, 3]);

  assert.deepEqual(JSON.parse(await s.next()), { type: "complete", ticks: 3 });
  const e = await s.closedWithin(3000);
  assert.equal(e.code, 1000);
  assert.equal(e.reason, "Ticker finished");
});

// The browser API only reports 1006 for a failed handshake, so check the HTTP answer directly.
test("ticker rejects a count that isn't a number at the handshake", async () => {
  const target = new URL(url("/ws/ticker?count=abc").replace(/^ws/, "http"));
  const res = await new Promise((resolve, reject) => {
    const req = request(target, {
      headers: {
        Connection: "Upgrade",
        Upgrade: "websocket",
        "Sec-WebSocket-Version": "13",
        "Sec-WebSocket-Key": "dGhlIHNhbXBsZSBub25jZQ==",
      },
    });
    req.on("response", resolve);
    req.on("upgrade", () => reject(new Error("the handshake should have failed")));
    req.on("error", reject);
    req.end();
  });

  assert.equal(res.statusCode, 400);
  let body = "";
  for await (const chunk of res) body += chunk;
  assert.deepEqual(JSON.parse(body).errors.count, ["The value 'abc' is not valid."]);
});

# test_websocket.py
import asyncio
import json
import os

import pytest
from websockets.asyncio.client import connect
from websockets.exceptions import ConnectionClosed, InvalidStatus

BASE = os.environ.get("WS_BASE", "wss://api.testingapis.com")
KEY = os.environ.get("TESTINGAPIS_KEY")


def url(path):
    if not KEY:
        return BASE + path
    return f"{BASE}{path}{'&' if '?' in path else '?'}ta_key={KEY}"


async def next_message(ws, timeout=5):
    return await asyncio.wait_for(ws.recv(), timeout)


@pytest.mark.asyncio
async def test_echo_returns_text_and_binary_unchanged():
    async with connect(url("/ws/echo")) as ws:
        welcome = json.loads(await next_message(ws))
        assert welcome["type"] == "welcome"

        await ws.send("héllo 👋")
        assert await next_message(ws) == "héllo 👋"

        await ws.send(bytes([0, 1, 254, 255]))
        assert await next_message(ws) == bytes([0, 1, 254, 255])


@pytest.mark.asyncio
async def test_sending_close_ends_with_1000_bye():
    async with connect(url("/ws/echo")) as ws:
        await next_message(ws)
        await ws.send("close")
        with pytest.raises(ConnectionClosed) as closed:
            await next_message(ws)

    assert closed.value.rcvd.code == 1000
    assert closed.value.rcvd.reason == "Bye!"


@pytest.mark.asyncio
async def test_ticker_sends_ticks_in_order_then_completes():
    async with connect(url("/ws/ticker?count=3")) as ws:
        ticks = [json.loads(await next_message(ws)) for _ in range(3)]
        assert [t["data"]["tick"] for t in ticks] == [1, 2, 3]
        assert json.loads(await next_message(ws)) == {"type": "complete", "ticks": 3}


@pytest.mark.asyncio
async def test_ticker_rejects_a_bad_count_at_the_handshake():
    # Unlike the browser API, this library reports the HTTP status of a failed handshake.
    with pytest.raises(InvalidStatus) as failed:
        async with connect(url("/ws/ticker?count=abc")):
            pass

    assert failed.value.response.status_code == 400