WebSocket

最近更新时间:2025-12-01 16:55:12

我的收藏
本文档提供了在云压测脚本模式中使用 WebSocket 协议进行压测的完整示例。所有示例均基于 pts/ws 模块提供的 WebSocket 连接和消息处理方法。

前置条件

在使用 WebSocket 功能前,请确保:
导入模块:在脚本开头使用 import ws from 'pts/ws' 导入 WebSocket 模块。
脚本结构:所有连接和消息处理代码需放在 export default function() 函数中执行。
URL 格式:确保目标 WebSocket URL 格式正确,支持 ws://wss:// 协议。
网络连通性:确保压测环境能够访问目标 WebSocket 服务器。
理解执行机制:WebSocket 脚本的执行机制与 HTTP 脚本不同:
HTTP 脚本的每个 VU(虚拟用户)会持续迭代主函数,直到压测结束。
WebSocket 脚本的每个 VU 不会持续迭代主函数,因为主函数会被 ws.connect 方法阻塞,直到连接关闭。
在连接建立后的回调函数中,会持续监听和处理异步事件,直到压测结束。

注意事项

连接生命周期:WebSocket 连接建立后,主函数会被阻塞,直到连接关闭。所有业务逻辑都应在 ws.connect() 的回调函数中实现。
事件监听顺序:应在连接建立后立即注册所有需要的事件监听器,确保不会遗漏任何消息。
消息发送时机:虽然可以在回调函数的任何位置调用 send(),但应在 open 事件触发后或确认连接已建立后再发送消息。
资源清理:使用 setTimeoutsetIntervalsetLoop 时,注意在连接关闭时清理这些定时器,避免资源泄漏。
消息格式:根据服务器要求选择正确的消息格式(文本或二进制),并使用相应的方法发送。
连接保持:对于长连接场景,建议定期发送 ping 消息或心跳消息,防止连接因超时被关闭。
并发控制:在压测场景中,每个 VU 会建立独立的 WebSocket 连接,注意控制并发连接数,避免对服务器造成过大压力。
状态码检查:连接建立后,应检查 res.status 是否为101,确认连接成功建立。

WebSocket 连接

函数说明

ws.connect() 用于建立 WebSocket 连接,是 WebSocket 压测的基础方法。该方法在脚本模式的压测场景中使用,支持建立持久连接并进行双向通信。连接建立后,主函数会被阻塞,直到连接关闭。

函数签名

ws.connect(url: string, callback: (socket: Socket) => void, headers?: Record<string, string>): Response

参数说明

url (string,必填):目标 WebSocket 服务器的完整 URL 地址,支持 ws://wss:// 协议。
callback (function,必填):连接建立成功后的回调函数,接收一个 Socket 对象作为参数。
回调函数会在连接建立后立即执行。
在回调函数中定义所有的消息发送和事件监听逻辑。
回调函数执行完成后,ws.connect 才会返回。
headers (Record<string, string>,可选):请求连接时的 headers 配置,用于设置自定义请求头。

返回值说明

返回 Response 对象,包含以下属性:
status (number):HTTP 状态码,WebSocket 连接成功时状态码为 101(Switching Protocols)。
body (string):响应包体内容。
headers (Record<string, string>):响应头参数。
url (string):请求地址。

使用限制

URL 必须包含完整的协议前缀(ws://wss://)。
回调函数是必需的,所有 Socket 操作都必须在回调函数中进行。
连接建立后,主函数会被阻塞,直到连接关闭或回调函数执行完成。
连接失败时,status 不是 101,需要检查连接状态。

示例

示例1:基本连接和消息收发

本示例演示如何建立 WebSocket 连接、发送消息和接收消息,适用于基本的 WebSocket 通信场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
// 建立 WebSocket 连接
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
// 监听连接建立事件
socket.on('open', () => {
console.log('WebSocket connected');
});
// 监听接收文本消息事件
socket.on('message', (data) => {
console.log('Message received:', data);
});
// 监听连接关闭事件
socket.on('close', () => {
console.log('WebSocket disconnected');
});
// 发送文本消息
socket.send('Hello, WebSocket!');
});
// 检查连接状态码是否为 101(WebSocket 协议切换成功)
check('status is 101', () => res.status === 101);
}

示例2:带错误处理的连接

本示例演示如何处理连接错误,适用于需要健壮错误处理的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected successfully');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
socket.on('close', () => {
console.log('Connection closed');
});
socket.send('Test message');
});
// 验证连接成功
check('connection successful', () => res.status === 101);
// 连接失败时输出错误信息
if (res.status !== 101) {
console.error('Connection failed with status:', res.status);
}
}

示例3:带自定义请求头的连接

本示例演示如何在建立连接时设置自定义请求头,适用于需要传递认证信息或自定义头部字段的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
// 定义自定义请求头
const headers = {
'X-MyApplication': 'PTS',
'X-MyScript': 'WebSocket',
'Authorization': 'Bearer token123'
};
// 建立连接并传入 headers
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected with custom headers');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
socket.on('close', () => {
console.log('Disconnected');
});
socket.send('Hello');
}, headers);
check('connection successful', () => res.status === 101);
}

Socket 对象事件监听

函数说明

socket.on() 用于监听 WebSocket 连接的各种事件,是处理异步消息和连接状态变化的核心方法。该方法在 ws.connect() 的回调函数中使用,用于注册事件处理函数。

函数签名

socket.on(event: string, callback: ((...args: any[]) => void)): void

参数说明

event (string,必填):事件名称,支持以下事件:
open:建立连接时触发,无参数。
close:关闭连接时触发,无参数。
message:接收文本消息时触发,回调函数接收消息内容作为参数。
binaryMessage:接收二进制消息时触发,回调函数接收 ArrayBuffer 作为参数。
ping:接收 ping 消息时触发,无参数。
pong:接收 pong 消息时触发,无参数。
error:发生错误时触发,回调函数接收错误对象作为参数,错误对象包含 error() 方法用于获取错误信息。
callback (function,必填):事件回调函数,根据事件类型接收不同的参数。

返回值说明

无返回值(void)。

使用限制

必须在 ws.connect() 的回调函数中调用。
同一个事件可以注册多个监听器,都会被执行。
事件监听器应该在发送消息之前注册,以确保能接收到所有消息。
messagebinaryMessage 事件回调函数的参数类型不同,需要根据消息类型选择正确的事件。

示例

示例1:监听所有事件类型

本示例演示如何监听 WebSocket 支持的所有事件类型,适用于需要全面监控连接状态的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
// 监听连接建立
socket.on('open', () => {
console.log('Connection opened');
});
// 监听文本消息
socket.on('message', (data) => {
console.log('Text message received:', data);
});
// 监听二进制消息
socket.on('binaryMessage', (data) => {
console.log('Binary message received, length:', data.byteLength);
});
// 监听 ping 消息
socket.on('ping', () => {
console.log('Ping received');
});
// 监听 pong 消息
socket.on('pong', () => {
console.log('Pong received');
});
// 监听连接关闭
socket.on('close', () => {
console.log('Connection closed');
});
// 监听错误事件
socket.on('error', (e) => {
console.log('Error happened:', e.error());
});
// 发送测试消息
socket.send('Hello');
});
check('connection successful', () => res.status === 101);
}

示例2:处理不同类型的消息

本示例演示如何区分处理文本消息和二进制消息,适用于需要处理多种消息格式的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
// 发送文本消息
socket.send('Text message');
// 发送二进制消息
socket.sendBinary(new ArrayBuffer(8));
});
// 处理文本消息
socket.on('message', (data) => {
console.log('Received text:', data);
// 可以在这里进行文本消息的业务处理
if (data === 'ping') {
socket.send('pong');
}
});
// 处理二进制消息
socket.on('binaryMessage', (data) => {
console.log('Received binary, size:', data.byteLength);
// 可以在这里进行二进制消息的业务处理
const view = new Uint8Array(data);
console.log('First byte:', view[0]);
});
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

发送文本消息

函数说明

socket.send() 用于向 WebSocket 服务器发送文本消息,是 WebSocket 双向通信中客户端向服务器发送数据的主要方法。该方法在连接建立后的回调函数中使用。

函数签名

socket.send(msg: string): void

参数说明

msg (string,必填):要发送的文本消息内容。

返回值说明

无返回值(void)。

使用限制

必须在连接建立后(open 事件触发后)调用,否则消息无法发送。
消息内容必须是字符串类型,非字符串值会被转换为字符串。
连接已关闭时,调用此方法不会报错,但消息不会发送。
应在 open 事件回调中或确认连接已建立后再发送消息。

示例

示例1:发送简单文本消息

本示例演示如何发送简单的文本消息,适用于基本的消息发送场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
// 连接建立后立即发送消息
socket.send('Hello, Server!');
});
socket.on('message', (data) => {
console.log('Echo received:', data);
});
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

示例2:发送 JSON 格式消息

本示例演示如何发送 JSON 格式的文本消息,适用于需要发送结构化数据的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
// 构造 JSON 对象并转换为字符串发送
const message = {
type: 'chat',
user: 'testuser',
content: 'Hello from WebSocket',
timestamp: Date.now()
};
socket.send(JSON.stringify(message));
});
socket.on('message', (data) => {
try {
// 解析接收到的 JSON 消息
const msg = JSON.parse(data);
console.log('Received message:', msg);
} catch (e) {
console.log('Received non-JSON message:', data);
}
});
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

示例3:定时发送消息

本示例演示如何定时发送消息,适用于需要周期性发送数据的场景。
import ws from 'pts/ws';
import { check, sleep } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
let messageCount = 0;
socket.on('open', () => {
console.log('Connected, starting to send messages');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
// 使用 setInterval 定时发送消息
socket.setInterval(function () {
messageCount++;
const msg = `Message #${messageCount}`;
socket.send(msg);
console.log('Sent:', msg);
}, 1000); // 每秒发送一条消息
// 10 秒后关闭连接
socket.setTimeout(function () {
console.log('Closing connection after 10 seconds');
socket.close();
}, 10000);
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

发送二进制消息

函数说明

socket.sendBinary() 用于向 WebSocket 服务器发送二进制消息,适用于需要传输二进制数据(如图片、文件、协议数据等)的场景。该方法在连接建立后的回调函数中使用。

函数签名

socket.sendBinary(msg: ArrayBuffer): void

参数说明

msg (ArrayBuffer,必填):要发送的二进制数据,必须是 ArrayBuffer 类型。

返回值说明

无返回值(void)。

使用限制

必须在连接建立后调用。
参数必须是 ArrayBuffer 类型,不能是其他类型(如 Uint8Array、Buffer 等)。
从其他格式转换为 ArrayBuffer 时,使用 util.toArrayBuffer() 进行转换。
二进制消息的接收需要使用 binaryMessage 事件,而不是 message 事件。

示例

示例1:发送简单二进制数据

本示例演示如何发送简单的二进制数据,适用于基本的二进制消息发送场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
// 创建一个包含 8 字节的 ArrayBuffer
const buffer = new ArrayBuffer(8);
const view = new Uint8Array(buffer);
// 填充一些测试数据
for (let i = 0; i < view.length; i++) {
view[i] = i;
}
// 发送二进制数据
socket.sendBinary(buffer);
console.log('Binary message sent, size:', buffer.byteLength);
});
// 监听二进制消息
socket.on('binaryMessage', (data) => {
console.log('Binary message received, size:', data.byteLength);
const view = new Uint8Array(data);
console.log('First 4 bytes:', view.slice(0, 4));
});
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

示例2:发送字符串转二进制

本示例演示如何将字符串转换为二进制数据发送,适用于需要以二进制格式传输文本的场景。
import ws from 'pts/ws';
import { check } from 'pts';
import util from 'pts/util';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
// 使用 PTS 的 util.toArrayBuffer 将字符串转换为 ArrayBuffer
const text = 'Hello, Binary World!';
const buffer = util.toArrayBuffer(text);
socket.sendBinary(buffer);
console.log('Sent binary data from string:', text);
});
socket.on('binaryMessage', (data) => {
// 读取二进制数据的字节信息
console.log('Received binary message, size:', data.byteLength);
// 使用 Uint8Array 读取字节
const view = new Uint8Array(data);
console.log('First 5 bytes:', Array.from(view.slice(0, 5)));
// 将 ArrayBuffer 转换为字符串:使用 base64 编码/解码
// 注意:仅适用于 ArrayBuffer 内容是 UTF-8 编码文本的场景
const base64 = util.base64Encoding(data);
const text = util.base64Decoding(base64);
console.log('Decoded text:', text);
});
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

关闭连接

函数说明

socket.close() 用于主动关闭 WebSocket 连接,适用于需要控制连接生命周期的场景。调用此方法后,连接会正常关闭,并触发 close 事件。

函数签名

socket.close(): void

返回值说明

无返回值(void)。

使用限制

可以在连接建立后的任何时候调用
调用后连接会立即关闭,无法再发送或接收消息
关闭连接会触发 close 事件
连接已经关闭时,再次调用此方法不会报错,但不会有任何效果

示例

示例1:定时关闭连接

本示例演示如何在指定时间后关闭连接,适用于需要限制连接时长的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
socket.send('Hello');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
// 5 秒后关闭连接
socket.setTimeout(function () {
console.log('5 seconds passed, closing connection');
socket.close();
}, 5000);
socket.on('close', () => {
console.log('Connection closed');
});
});
check('connection successful', () => res.status === 101);
}

示例2:收到特定消息后关闭

本示例演示如何在收到特定消息后关闭连接,适用于需要根据服务器响应决定是否关闭连接的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
socket.send('start');
});
socket.on('message', (data) => {
console.log('Received:', data);
// 收到 "close" 消息时关闭连接
if (data === 'close' || data === 'bye') {
console.log('Received close signal, closing connection');
socket.close();
} else {
// 继续发送消息
socket.send('continue');
}
});
socket.on('close', () => {
console.log('Connection closed');
});
});
check('connection successful', () => res.status === 101);
}

发送 Ping 消息

函数说明

socket.ping() 用于向服务器发送 WebSocket ping 消息,用于保持连接活跃和检测连接状态。服务器通常会响应 pong 消息。该方法在连接建立后的回调函数中使用。

函数签名

socket.ping(): void

返回值说明

无返回值(void)。

使用限制

必须在连接建立后调用
ping 消息是 WebSocket 协议层面的控制消息,不会触发 message 事件
服务器响应 pong 消息时,会触发 pong 事件
定期发送 ping 可以保持连接活跃,防止因超时被关闭

示例:定期发送 Ping 保持连接

本示例演示如何定期发送 ping 消息来保持连接活跃,适用于需要维持长连接的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
});
// 监听 pong 响应
socket.on('pong', () => {
console.log('Received pong from server');
});
// 每 500 毫秒发送一次 ping
socket.setInterval(function () {
socket.ping();
console.log('Sent ping');
}, 500);
socket.on('message', (data) => {
console.log('Received message:', data);
});
// 30 秒后关闭连接
socket.setTimeout(function () {
console.log('Closing connection');
socket.close();
}, 30000);
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

设置定时器

函数说明

socket.setTimeout() 用于在 WebSocket 连接上设置一个定时器,在指定时间后执行回调函数。适用于需要在特定时间后执行某些操作(如关闭连接、发送消息等)的场景。

函数签名

socket.setTimeout(callback: (() => void), intervalMs: number): void

参数说明

callback (function,必填):定时器到期后执行的回调函数,无参数。
intervalMs (number,必填):定时器延迟时间,单位为毫秒。

返回值说明

无返回值(void)。

使用限制

必须在连接建立后的回调函数中调用。
定时器是单次执行的,执行一次后不会自动重复。
需要重复执行时,可以在回调函数中再次调用 setTimeout,或使用 setInterval
连接在定时器到期前关闭时,定时器不会执行。

示例

示例1:延迟发送消息

本示例演示如何使用定时器延迟发送消息,适用于需要延迟操作的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
// 2 秒后发送第一条消息
socket.setTimeout(function () {
socket.send('First message after 2 seconds');
}, 2000);
// 5 秒后发送第二条消息
socket.setTimeout(function () {
socket.send('Second message after 5 seconds');
}, 5000);
});
socket.on('message', (data) => {
console.log('Received:', data);
});
// 10 秒后关闭连接
socket.setTimeout(function () {
console.log('Closing connection');
socket.close();
}, 10000);
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

示例2:超时自动关闭

本示例演示如何使用定时器实现超时自动关闭连接,适用于需要限制连接时长的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
let isClosed = false;
socket.on('open', () => {
console.log('Connected, will close after 5 seconds');
socket.send('Hello');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
// 5 秒后自动关闭连接
socket.setTimeout(function () {
if (!isClosed) {
console.log('Timeout reached, closing connection');
isClosed = true;
socket.close();
}
}, 5000);
socket.on('close', () => {
console.log('Connection closed');
isClosed = true;
});
});
check('connection successful', () => res.status === 101);
}

设置轮询函数

函数说明

socket.setInterval() 用于在 WebSocket 连接上设置一个轮询函数,按照指定的时间间隔重复执行。适用于需要周期性执行某些操作(如定期发送消息、定期 ping 等)的场景。

函数签名

socket.setInterval(callback: (() => void), intervalMs: number): void

参数说明

callback (function,必填):每次轮询时执行的回调函数,无参数。
intervalMs (number,必填):轮询间隔时间,单位为毫秒。

返回值说明

无返回值(void)。

使用限制

必须在连接建立后的回调函数中调用
轮询会持续执行,直到连接关闭
回调函数执行时间超过间隔时间时,会影响定时精度
应在回调函数中避免耗时操作,或使用异步处理

示例

示例1:定期发送消息

本示例演示如何使用轮询函数定期发送消息,适用于需要周期性发送数据的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
let messageCount = 0;
socket.on('open', () => {
console.log('Connected, starting periodic messages');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
// 每 1 秒发送一条消息
socket.setInterval(function () {
messageCount++;
const msg = `Periodic message #${messageCount}`;
socket.send(msg);
console.log('Sent:', msg);
}, 1000);
// 10 秒后关闭连接
socket.setTimeout(function () {
console.log('Closing connection');
socket.close();
}, 10000);
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

示例2:定期发送心跳

本示例演示如何使用轮询函数定期发送心跳消息,适用于需要保持连接活跃的场景。
import ws from 'pts/ws';
import { check } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
socket.on('open', () => {
console.log('Connected');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
// 每 2 秒发送一次心跳消息
socket.setInterval(function () {
const heartbeat = {
type: 'heartbeat',
timestamp: Date.now()
};
socket.send(JSON.stringify(heartbeat));
console.log('Sent heartbeat');
}, 2000);
// 30 秒后关闭连接
socket.setTimeout(function () {
console.log('Closing connection');
socket.close();
}, 30000);
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

设置循环执行函数

函数说明

socket.setLoop() 用于在 WebSocket 连接上设置一个循环执行函数,该函数会持续执行直到连接关闭。与 setInterval 不同,setLoop 的执行频率由函数内部的 sleep() 调用控制,提供了更灵活的执行控制。

函数签名

socket.setLoop(callback: (() => void)): void

参数说明

callback (function,必填):循环执行的回调函数,无参数。函数内部应包含 sleep() 调用来控制执行频率。

返回值说明

无返回值(void)。

使用限制

必须在连接建立后的回调函数中调用
循环会持续执行,直到连接关闭
回调函数内部必须调用 sleep() 来控制执行频率,否则会无限快速执行
建议在回调函数中使用 sleep() 来避免 CPU 占用过高

示例

示例 1:循环发送消息

本示例演示如何使用循环函数持续发送消息,适用于需要持续发送数据的场景。
import ws from 'pts/ws';
import { check, sleep } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
let messageCount = 0;
socket.on('open', () => {
console.log('Connected, starting loop');
});
socket.on('message', (data) => {
console.log('Received:', data);
});
// 设置循环,每 0.5 秒发送一条消息
socket.setLoop(function () {
messageCount++;
socket.send(`Loop message #${messageCount}`);
console.log('Sent loop message #' + messageCount);
// 必须调用 sleep 来控制执行频率
sleep(0.5);
});
// 10 秒后关闭连接
socket.setTimeout(function () {
console.log('Closing connection');
socket.close();
}, 10000);
socket.on('close', () => {
console.log('Disconnected');
});
});
check('connection successful', () => res.status === 101);
}

示例 2:循环处理业务逻辑

本示例演示如何在循环中处理复杂的业务逻辑,适用于需要持续处理业务的场景。
import ws from 'pts/ws';
import { check, sleep } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
let state = 'idle';
socket.on('open', () => {
console.log('Connected');
state = 'active';
});
socket.on('message', (data) => {
console.log('Received:', data);
// 根据接收到的消息更新状态
if (data === 'start') {
state = 'running';
} else if (data === 'stop') {
state = 'stopped';
}
});
// 循环处理业务逻辑
socket.setLoop(function () {
if (state === 'running') {
// 业务逻辑:发送数据
const data = {
action: 'process',
timestamp: Date.now()
};
socket.send(JSON.stringify(data));
console.log('Sent process data');
} else if (state === 'stopped') {
console.log('State is stopped, waiting...');
}
// 控制循环频率:每 1 秒执行一次
sleep(1);
});
// 30 秒后关闭连接
socket.setTimeout(function () {
console.log('Closing connection');
socket.close();
}, 30000);
socket.on('close', () => {
console.log('Disconnected');
state = 'closed';
});
});
check('connection successful', () => res.status === 101);
}

综合示例:完整的 WebSocket 压测场景

本示例演示一个完整的 WebSocket 压测场景,包含连接建立、消息收发、心跳保持、错误处理和连接关闭等完整流程。
import ws from 'pts/ws';
import { check, sleep } from 'pts';

export default function () {
const res = ws.connect('ws://mockwebsocket.pts.svc.cluster.local/echo', function (socket) {
let messageCount = 0;
let isActive = true;
// 监听连接建立
socket.on('open', () => {
console.log('WebSocket connection established');
check('connection opened', () => true);
});
// 监听文本消息
socket.on('message', (data) => {
messageCount++;
console.log(`Message #${messageCount} received:`, data);
// 验证消息内容
if (data.includes('echo')) {
check('message contains echo', () => data.includes('echo'));
}
});
// 监听二进制消息
socket.on('binaryMessage', (data) => {
console.log('Binary message received, size:', data.byteLength);
});
// 监听 pong 响应
socket.on('pong', () => {
console.log('Pong received from server');
});
// 监听连接关闭
socket.on('close', () => {
console.log('WebSocket connection closed');
isActive = false;
});
// 连接建立后立即发送欢迎消息
socket.send('Hello from WebSocket client');
// 定期发送 ping 保持连接
socket.setInterval(function () {
if (isActive) {
socket.ping();
console.log('Sent ping');
}
}, 2000);
// 循环发送业务消息
socket.setLoop(function () {
if (isActive) {
const message = {
type: 'data',
sequence: messageCount + 1,
timestamp: Date.now(),
data: `Message content ${messageCount + 1}`
};
socket.send(JSON.stringify(message));
console.log('Sent business message:', message.sequence);
}
sleep(1); // 每 1 秒发送一次
});
// 30 秒后关闭连接
socket.setTimeout(function () {
if (isActive) {
console.log('Closing connection after 30 seconds');
socket.close();
}
}, 30000);
});
// 验证连接状态
check('WebSocket connection successful', () => res.status === 101);
if (res.status !== 101) {
console.error('WebSocket connection failed with status:', res.status);
}
}