appium/packages/base-driver/lib/express/websocket.ts

import {util} from '@appium/support';
import type {AppiumServer, WSServer} from '@appium/types';

export const DEFAULT_WS_PATHNAME_PREFIX = '/ws';

/**
 * Adds a WebSocket handler to this server's mapping.
 * @see AppiumServerExtension.addWebSocketHandler
 */
export async function addWebSocketHandler(
  this: AppiumServer,
  handlerPathname: string,
  handlerServer: WSServer,
): Promise<void> {
  this.webSocketsMapping[handlerPathname] = handlerServer;
}

/**
 * Returns WebSocket handlers for this server, optionally filtered by pathname.
 * @see AppiumServerExtension.getWebSocketHandlers
 */
export async function getWebSocketHandlers(
  this: AppiumServer,
  keysFilter: string | null = null,
): Promise<Record<string, WSServer>> {
  return Object.entries(this.webSocketsMapping).reduce<Record<string, WSServer>>((acc, [pathname, wsServer]) => {
    if (typeof keysFilter !== 'string' || pathname.includes(keysFilter)) {
      acc[pathname] = wsServer;
    }
    return acc;
  }, {});
}

/**
 * Removes a WebSocket handler by pathname.
 * @see AppiumServerExtension.removeWebSocketHandler
 */
export async function removeWebSocketHandler(this: AppiumServer, handlerPathname: string): Promise<boolean> {
  const wsServer = this.webSocketsMapping?.[handlerPathname];
  if (!wsServer) {
    return false;
  }

  try {
    wsServer.close();
    for (const client of wsServer.clients || []) {
      client.terminate();
    }
    return true;
  } catch {
    // ignore
  } finally {
    delete this.webSocketsMapping[handlerPathname];
  }
  return false;
}

/**
 * Removes all WebSocket handlers from this server.
 * @see AppiumServerExtension.removeAllWebSocketHandlers
 */
export async function removeAllWebSocketHandlers(this: AppiumServer): Promise<boolean> {
  if (util.isEmpty(this.webSocketsMapping)) {
    return false;
  }

  const results = await Promise.all(
    Object.keys(this.webSocketsMapping).map((pathname) => this.removeWebSocketHandler(pathname)),
  );
  return results.some(Boolean);
}