WebSocket

基础示例

ws.mock.ts

import type { MockRequest } from 'vite-plugin-mock-dev-server'
import { Buffer } from 'node:buffer'
import { defineMock } from 'vite-plugin-mock-dev-server'

/**
 * Convert the raw websocket data to text
 *
 * 将 websocket 原始数据转换为文本
 *
 * @param raw - Raw data / 原始数据
 * @returns Message text / 消息文本
 */
function rawToString(raw: unknown): string {
  if (Buffer.isBuffer(raw)) {
    return raw.toString()
  }
  if (Array.isArray(raw)) {
    return Buffer.concat(raw).toString()
  }
  if (raw instanceof ArrayBuffer) {
    return Buffer.from(raw).toString()
  }
  return typeof raw === 'string' ? raw : ''
}

/**
 * WebSocket mocks: connection, message echo / broadcast, dynamic params, query,
 * cookies, server push and `onCleanup`.
 *
 * WebSocket mock：连接、消息回显 / 广播、动态参数、query、cookie、服务端推送，
 * 以及 `onCleanup` 清理。
 */
export default defineMock([
  {
    url: '/socket.io',
    ws: true,
    setup(wss, { onCleanup }) {
      const heartbeat = setInterval(() => {
        wss.clients.forEach((client) => {
          if (client.readyState === 1) {
            client.send(JSON.stringify({ type: 'heartbeat', at: Date.now() }))
          }
        })
      }, 5000)
      onCleanup(() => clearInterval(heartbeat))

      wss.on('connection', (ws, req: MockRequest) => {
        ws.send(
          JSON.stringify({
            type: 'connected',
            url: req.url,
            query: req.query,
            params: req.params,
            cookie: req.getCookie('token') ?? null,
          }),
        )

        ws.on('message', (raw) => {
          const text = rawToString(raw)
          let message: any
          try {
            message = JSON.parse(text)
          } catch {
            message = { type: 'text', payload: text }
          }

          if (message?.type === 'broadcast') {
            const data = JSON.stringify({
              type: 'broadcast',
              payload: message.payload,
              at: Date.now(),
            })
            wss.clients.forEach((client) => {
              if (client.readyState === 1) {
                client.send(data)
              }
            })
            return
          }

          ws.send(JSON.stringify({ type: 'echo', payload: message, at: Date.now() }))
        })
      })
    },
  },
  {
    url: '/socket.io/room/:room',
    ws: true,
    setup(wss, { onCleanup }) {
      const timer = setInterval(() => {
        wss.clients.forEach((client) => {
          if (client.readyState === 1) {
            client.send(JSON.stringify({ type: 'tick', at: Date.now() }))
          }
        })
      }, 2000)
      onCleanup(() => clearInterval(timer))

      wss.on('connection', (ws, req: MockRequest) => {
        ws.send(
          JSON.stringify({
            type: 'joined',
            room: req.params.room,
            url: req.url,
          }),
        )

        ws.on('message', (raw) => {
          const data = JSON.stringify({
            type: 'message',
            room: req.params.room,
            payload: rawToString(raw),
          })
          wss.clients.forEach((client) => {
            if (client.readyState === 1) {
              client.send(data)
            }
          })
        })
      })
    },
  },
])

模拟实时推送

ws-real-time-push.mock.ts

import { defineMock } from 'vite-plugin-mock-dev-server'

export default defineMock({
  url: '/socket/push',
  ws: true,
  setup(wss, { onCleanup }) {
    wss.on('connection', (ws, req) => {
      const timer1 = setInterval(() => {
        ws.send('push message')
      }, 1000)
      onCleanup(() => clearInterval(timer1))

      const timer2 = setInterval(() => {
        ws.send('push message2')
      }, 3000)
      onCleanup(() => clearInterval(timer2))
    })
  },
})

模拟聊天室

ws-chat.mock.ts

import type { MockRequest } from 'vite-plugin-mock-dev-server'
import { defineMock } from 'vite-plugin-mock-dev-server'

export default defineMock({
  url: '/socket/chat',
  ws: true,
  setup(wss) {
    const userList = []
    const userWsMap = {}
    wss.on('connection', (ws, req: MockRequest) => {
      const token = req.getCookie('token')
      const user = getUserByToken(token)

      userWsMap[user.id] = ws
      if (userList.findIndex(({ id }) => id === user.Id) !== -1) {
        userList.push(user)
        // 将当前成员更新广播给聊天室所有人
        userList.forEach(({ id }) =>
          userWsMap[id].send(
            JSON.stringify({
              type: 'userList',
              data: { userList },
            }),
          ),
        )
      }
      ws.on('message', (raw) => {
        const {
          type,
          data: { userId, message },
        } = JSON.parse(String(raw))
        // 服务器接收用户发送到的聊天信息，将其广播给聊天室中的其他成员
        if (type === 'chat') {
          userList.forEach((user) => {
            user.id !== userId &&
              userWsMap[user.id].send(JSON.stringify({ type: 'chat', data: { userId, message } }))
          })
        }
      })

      ws.on('close', () => {
        // 有成员退出
        const idx = userList.findIndex(({ id }) => user.id === id)
        idx !== -1 && userList.splice(idx, 1)
        userWsMap[user.id] = null
        // 将当前成员更新广播给聊天室所有人
        userList.forEach(({ id }) =>
          userWsMap[id].send(
            JSON.stringify({
              type: 'userList',
              data: { userList },
            }),
          ),
        )
      })
    })
  },
})

app.ts

const store = {
  userList: [],
  chatList: [],
  connected: false,
  currentUserId: '112233',
}
const ws = new WebSocket('ws://localhost:5173/socket/chat')
ws.addEventListener(
  'open',
  () => {
    store.connected = true
    // heartbeat
    setInterval(() => {
      ws.send(JSON.stringify({ type: 'heartbeat' }))
    }, 1000)
  },
  { once: true },
)
ws.addEventListener('message', (raw) => {
  const { type, data } = JSON.parse(String(raw))
  if (type === 'userList') store.userList = data.userList || []

  if (type === 'chat') store.chatList.push(data)
})

function sendMessage(message: string) {
  if (store.connected) {
    ws.send(
      JSON.stringify({
        type: 'chat',
        data: { userId: store.currentUserId, message },
      }),
    )
  }
}
function closeChat() {
  if (store.connected) ws.close()
}