一、前置条件
在使用 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 事件触发后或确认连接已建立后再发送消息。
·资源清理：使用 setTimeout、setInterval 和 setLoop 时，注意在连接关闭时清理这些定时器，避免资源泄漏。
·消息格式：根据服务器要求选择正确的消息格式（文本或二进制），并使用相应的方法发送。
·连接保持：对于长连接场景，建议定期发送 ping 消息或心跳消息，防止连接因超时被关闭。
·并发控制：在压测场景中，每个 VU 会建立独立的 WebSocket 连接，注意控制并发连接数，避免对服务器造成过大压力。
·状态码检查：连接建立后，应检查 res.status 是否为101，确认连接成功建立。
四、函数说明
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.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() 的回调函数中调用。
·同一个事件可以注册多个监听器，都会被执行。
·事件监听器应该在发送消息之前注册，以确保能接收到所有消息。
·message 和 binaryMessage 事件回调函数的参数类型不同，需要根据消息类型选择正确的事件。
20、示例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);}
21、示例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);}
23、函数说明
socket.send() 用于向 WebSocket 服务器发送文本消息，是 WebSocket 双向通信中客户端向服务器发送数据的主要方法。该方法在连接建立后的回调函数中使用。
24、函数签名
socket.send(msg: string): void
25、参数说明
msg (string，必填)：要发送的文本消息内容。
26、返回值说明
无返回值（void）。
27、使用限制
·必须在连接建立后（open 事件触发后）调用，否则消息无法发送。
·消息内容必须是字符串类型，非字符串值会被转换为字符串。
·连接已关闭时，调用此方法不会报错，但消息不会发送。
·应在 open 事件回调中或确认连接已建立后再发送消息。
29、示例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);}
30、示例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);}
31、示例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);}
33、函数说明
socket.sendBinary() 用于向 WebSocket 服务器发送二进制消息，适用于需要传输二进制数据（如图片、文件、协议数据等）的场景。该方法在连接建立后的回调函数中使用。
34、函数签名
socket.sendBinary(msg: ArrayBuffer): void
35、参数说明
msg (ArrayBuffer，必填)：要发送的二进制数据，必须是 ArrayBuffer 类型。
36、返回值说明
无返回值（void）。
37、使用限制
·必须在连接建立后调用。
·参数必须是 ArrayBuffer 类型，不能是其他类型（如 Uint8Array、Buffer 等）。
·从其他格式转换为 ArrayBuffer 时，使用 util.toArrayBuffer() 进行转换。
·二进制消息的接收需要使用 binaryMessage 事件，而不是 message 事件。
39、示例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);}
40、示例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);}
42、函数说明
socket.close() 用于主动关闭 WebSocket 连接，适用于需要控制连接生命周期的场景。调用此方法后，连接会正常关闭，并触发 close 事件。
43、函数签名
socket.close(): void
44、返回值说明
无返回值（void）。
45、使用限制
·可以在连接建立后的任何时候调用
·调用后连接会立即关闭，无法再发送或接收消息
·关闭连接会触发 close 事件
·连接已经关闭时，再次调用此方法不会报错，但不会有任何效果
47、示例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);}
48、示例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);}
50、函数说明
socket.ping() 用于向服务器发送 WebSocket ping 消息，用于保持连接活跃和检测连接状态。服务器通常会响应 pong 消息。该方法在连接建立后的回调函数中使用。
51、函数签名
socket.ping(): void
52、返回值说明
无返回值（void）。
53、使用限制
·必须在连接建立后调用
·ping 消息是 WebSocket 协议层面的控制消息，不会触发 message 事件
·服务器响应 pong 消息时，会触发 pong 事件
·定期发送 ping 可以保持连接活跃，防止因超时被关闭
54、示例：定期发送 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);}
56、函数说明
socket.setTimeout() 用于在 WebSocket 连接上设置一个定时器，在指定时间后执行回调函数。适用于需要在特定时间后执行某些操作（如关闭连接、发送消息等）的场景。
57、函数签名
socket.setTimeout(callback: (() => void), intervalMs: number): void
58、参数说明
·callback (function，必填)：定时器到期后执行的回调函数，无参数。
·intervalMs (number，必填)：定时器延迟时间，单位为毫秒。
59、返回值说明
无返回值（void）。
60、使用限制
·必须在连接建立后的回调函数中调用。
·定时器是单次执行的，执行一次后不会自动重复。
·需要重复执行时，可以在回调函数中再次调用 setTimeout，或使用 setInterval。
·连接在定时器到期前关闭时，定时器不会执行。
62、示例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);}
63、示例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);}
65、函数说明
socket.setInterval() 用于在 WebSocket 连接上设置一个轮询函数，按照指定的时间间隔重复执行。适用于需要周期性执行某些操作（如定期发送消息、定期 ping 等）的场景。
66、函数签名
socket.setInterval(callback: (() => void), intervalMs: number): void
67、参数说明
·callback (function，必填)：每次轮询时执行的回调函数，无参数。
·intervalMs (number，必填)：轮询间隔时间，单位为毫秒。
68、返回值说明
无返回值（void）。
69、使用限制
·必须在连接建立后的回调函数中调用
·轮询会持续执行，直到连接关闭
·回调函数执行时间超过间隔时间时，会影响定时精度
·应在回调函数中避免耗时操作，或使用异步处理
71、示例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);}
72、示例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);}
74、函数说明
socket.setLoop() 用于在 WebSocket 连接上设置一个循环执行函数，该函数会持续执行直到连接关闭。与 setInterval 不同，setLoop 的执行频率由函数内部的 sleep() 调用控制，提供了更灵活的执行控制。
75、函数签名
socket.setLoop(callback: (() => void)): void
76、参数说明
callback (function，必填)：循环执行的回调函数，无参数。函数内部应包含 sleep() 调用来控制执行频率。
77、返回值说明
无返回值（void）。
78、使用限制
·必须在连接建立后的回调函数中调用
·循环会持续执行，直到连接关闭
·回调函数内部必须调用 sleep() 来控制执行频率，否则会无限快速执行
·建议在回调函数中使用 sleep() 来避免 CPU 占用过高
80、示例 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);}
81、示例 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);}
82、综合示例：完整的 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); }}