@zimic/interceptor 1.2.3 → 1.2.4-canary.1

package/src/webSocket/WebSocketHandler.ts

@@ -12,7 +12,7 @@ import {
 } from '@/utils/webSocket';
 
 import { WEB_SOCKET_CONTROL_MESSAGES, WebSocketControlMessage } from './constants';
-import InvalidWebSocketMessage from './errors/InvalidWebSocketMessage';
+import InvalidWebSocketMessageError from './errors/InvalidWebSocketMessageError';
 import NotRunningWebSocketHandlerError from './errors/NotRunningWebSocketHandlerError';
 import {
   WebSocketEventMessageListener,
@@ -26,6 +26,10 @@ import {
   WebSocketMessage,
 } from './types';
 
+interface WebSocketRequestAbortOptions<Schema extends WebSocketSchema> {
+  shouldAbortRequest?: (request: WebSocketEventMessage<Schema>) => boolean;
+}
+
 abstract class WebSocketHandler<Schema extends WebSocketSchema> {
   private sockets = new Set<ClientSocket>();
 
@@ -40,7 +44,7 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
   } = {};
 
   private socketListeners = {
-    messageAbort: new Map<ClientSocket, Set<(error: unknown) => void>>(),
+    abortRequests: new Map<ClientSocket, Set<(options: WebSocketRequestAbortOptions<Schema>) => void>>(),
   };
 
   protected constructor(options: { socketTimeout?: number; messageTimeout?: number }) {
@@ -71,11 +75,14 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     socket.addEventListener('error', handleSocketError);
 
     const handleSocketClose = () => {
+      this.sockets.delete(socket);
+
+      this.emitSocket('abortRequests', socket);
+      this.socketListeners.abortRequests.delete(socket);
+
       socket.removeEventListener('message', handleSocketMessage);
       socket.removeEventListener('close', handleSocketClose);
       socket.removeEventListener('error', handleSocketError);
-
-      this.removeSocket(socket);
     };
 
     socket.addEventListener('close', handleSocketClose);
@@ -109,7 +116,7 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     if (typeof data === 'string') {
       return data;
     } else {
-      throw new InvalidWebSocketMessage(data);
+      throw new InvalidWebSocketMessageError(data);
     }
   }
 
@@ -119,11 +126,11 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     try {
       parsedMessage = JSON.parse(stringifiedMessage) as unknown;
     } catch {
-      throw new InvalidWebSocketMessage(stringifiedMessage);
+      throw new InvalidWebSocketMessageError(stringifiedMessage);
     }
 
     if (!this.isMessage(parsedMessage)) {
-      throw new InvalidWebSocketMessage(stringifiedMessage);
+      throw new InvalidWebSocketMessageError(stringifiedMessage);
     }
 
     if (this.isReplyMessage(parsedMessage)) {
@@ -154,6 +161,13 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     );
   }
 
+  isChannelEvent<Channel extends WebSocketChannel<Schema>>(
+    event: WebSocketEventMessage<Schema>,
+    channel: Channel,
+  ): event is WebSocketEventMessage<Schema, Channel> {
+    return event.channel === channel;
+  }
+
   private async notifyListeners(message: WebSocketMessage<Schema>, socket: ClientSocket) {
     if (this.isReplyMessage(message)) {
       await this.notifyReplyListeners(message, socket);
@@ -193,11 +207,6 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     await Promise.all(closingPromises);
   }
 
-  private removeSocket(socket: ClientSocket) {
-    this.abortSocketMessages([socket]);
-    this.sockets.delete(socket);
-  }
-
   private async createEventMessage<Channel extends WebSocketChannel<Schema>>(
     channel: Channel,
     eventData: WebSocketEventMessage<Schema, Channel>['data'],
@@ -233,43 +242,67 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     const request = await this.createEventMessage(channel, requestData);
     this.sendMessage(request, options.sockets);
 
-    const response = await this.waitForReply(channel, request.id, options.sockets);
+    const response = await this.waitForReply(channel, request, options.sockets);
     return response.data;
   }
 
   async waitForReply<Channel extends WebSocketChannelWithReply<Schema>>(
     channel: Channel,
-    requestId: WebSocketEventMessage<Schema, Channel>['id'],
+    request: WebSocketEventMessage<Schema, Channel>,
     sockets: Collection<ClientSocket> = this.sockets,
   ) {
     return new Promise<WebSocketReplyMessage<Schema, Channel>>((resolve, reject) => {
       const replyTimeout = setTimeout(() => {
-        this.offReply(channel, replyListener); // eslint-disable-line @typescript-eslint/no-use-before-define
-        this.offAbortSocketMessages(sockets, abortListener); // eslint-disable-line @typescript-eslint/no-use-before-define
+        this.offChannel('reply', channel, replyListener); // eslint-disable-line @typescript-eslint/no-use-before-define
+
+        for (const socket of sockets) {
+          this.offSocket('abortRequests', socket, abortRequestsHandler); // eslint-disable-line @typescript-eslint/no-use-before-define
+        }
 
         const timeoutError = new WebSocketMessageTimeoutError(this.messageTimeout);
         reject(timeoutError);
       }, this.messageTimeout);
 
-      const abortListener = this.onAbortSocketMessages(sockets, (error) => {
+      const replyListener = this.onChannel('reply', channel, (message) => {
+        if (message.requestId !== request.id) {
+          return;
+        }
+
         clearTimeout(replyTimeout);
 
-        this.offReply(channel, replyListener); // eslint-disable-line @typescript-eslint/no-use-before-define
-        this.offAbortSocketMessages(sockets, abortListener);
+        this.offChannel('reply', channel, replyListener);
+
+        for (const socket of sockets) {
+          this.offSocket('abortRequests', socket, abortRequestsHandler); // eslint-disable-line @typescript-eslint/no-use-before-define
+        }
 
-        reject(error);
+        resolve(message);
       });
 
-      const replyListener = this.onReply(channel, (message) => {
-        if (message.requestId === requestId) {
-          clearTimeout(replyTimeout);
+      const abortRequestsHandler = (options: WebSocketRequestAbortOptions<Schema>) => {
+        const shouldAbortRequest = options.shouldAbortRequest === undefined || options.shouldAbortRequest(request);
+
+        /* istanbul ignore if -- @preserve
+         * Aborting requests is highly non-deterministic because it depends on specific timing of socket events. */
+        if (!shouldAbortRequest) {
+          return;
+        }
+
+        clearTimeout(replyTimeout);
 
-          this.offReply(channel, replyListener);
-          this.offAbortSocketMessages(sockets, abortListener);
+        this.offChannel('reply', channel, replyListener);
 
-          resolve(message);
+        for (const socket of sockets) {
+          this.offSocket('abortRequests', socket, abortRequestsHandler);
         }
-      });
+
+        const abortError = new WebSocketMessageAbortError();
+        reject(abortError);
+      };
+
+      for (const socket of sockets) {
+        this.onSocket('abortRequests', socket, abortRequestsHandler);
+      }
     });
   }
 
@@ -322,12 +355,22 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     }
   }
 
-  onEvent<Channel extends WebSocketChannel<Schema>, Listener extends WebSocketEventMessageListener<Schema, Channel>>(
+  onChannel<Channel extends WebSocketChannel<Schema>, Listener extends WebSocketEventMessageListener<Schema, Channel>>(
+    type: 'event',
     channel: Channel,
-    listener: Listener,
-  ): Listener {
+    eventListener: Listener,
+  ): Listener;
+  onChannel<Channel extends WebSocketChannel<Schema>, Listener extends WebSocketReplyMessageListener<Schema, Channel>>(
+    type: 'reply',
+    channel: Channel,
+    replyListener: Listener,
+  ): Listener;
+  onChannel<
+    Channel extends WebSocketChannel<Schema>,
+    Listener extends WebSocketEventMessageListener<Schema, Channel> & WebSocketReplyMessageListener<Schema, Channel>,
+  >(type: 'event' | 'reply', channel: Channel, listener: Listener): Listener {
     const listeners = this.getOrCreateChannelListeners<Channel>(channel);
-    listeners.event.add(listener);
+    listeners[type].add(listener);
     return listener;
   }
 
@@ -344,59 +387,66 @@ abstract class WebSocketHandler<Schema extends WebSocketSchema> {
     return listeners;
   }
 
-  onReply<
-    Channel extends WebSocketChannelWithReply<Schema>,
-    Listener extends WebSocketReplyMessageListener<Schema, Channel>,
-  >(channel: Channel, listener: Listener): Listener {
-    const listeners = this.getOrCreateChannelListeners<Channel>(channel);
-    listeners.reply.add(listener);
-    return listener;
-  }
-
-  offEvent<Channel extends WebSocketChannel<Schema>>(
+  offChannel<Channel extends WebSocketChannel<Schema>>(
+    type: 'event',
     channel: Channel,
-    listener: WebSocketEventMessageListener<Schema, Channel>,
-  ) {
-    this.channelListeners[channel]?.event.delete(listener);
-  }
-
-  offReply<Channel extends WebSocketChannelWithReply<Schema>>(
+    eventListener: WebSocketEventMessageListener<Schema, Channel>,
+  ): void;
+  offChannel<Channel extends WebSocketChannel<Schema>>(
+    type: 'reply',
+    channel: Channel,
+    replyListener: WebSocketReplyMessageListener<Schema, Channel>,
+  ): void;
+  offChannel<Channel extends WebSocketChannel<Schema>>(
+    type: 'event' | 'reply',
     channel: Channel,
-    listener: WebSocketReplyMessageListener<Schema, Channel>,
+    listener: WebSocketEventMessageListener<Schema, Channel> & WebSocketReplyMessageListener<Schema, Channel>,
   ) {
-    this.channelListeners[channel]?.reply.delete(listener);
+    const listeners = this.channelListeners[channel];
+    listeners?.[type].delete(listener);
   }
 
-  removeAllChannelListeners() {
-    this.channelListeners = {};
+  onSocket<Listener extends (options: WebSocketRequestAbortOptions<Schema>) => void>(
+    type: 'abortRequests',
+    socket: ClientSocket,
+    listener: Listener,
+  ): Listener {
+    const listeners = this.getOrCreateSocketListeners(type, socket);
+    listeners.add(listener);
+    return listener;
   }
 
-  private onAbortSocketMessages(sockets: Collection<ClientSocket>, listener: (error: unknown) => void) {
-    for (const socket of sockets) {
-      let listeners = this.socketListeners.messageAbort.get(socket);
-      if (!listeners) {
-        listeners = new Set();
-        this.socketListeners.messageAbort.set(socket, listeners);
-      }
-      listeners.add(listener);
+  private getOrCreateSocketListeners(type: 'abortRequests', socket: ClientSocket) {
+    const listeners = this.socketListeners[type].get(socket) ?? new Set();
+
+    if (!this.socketListeners[type].has(socket)) {
+      this.socketListeners[type].set(socket, listeners);
     }
 
-    return listener;
+    return listeners;
   }
 
-  private offAbortSocketMessages(sockets: Collection<ClientSocket>, listener: (error: unknown) => void) {
-    for (const socket of sockets) {
-      this.socketListeners.messageAbort.get(socket)?.delete(listener);
+  offSocket<Listener extends (options: WebSocketRequestAbortOptions<Schema>) => void>(
+    type: 'abortRequests',
+    socket: ClientSocket,
+    listener: Listener,
+  ) {
+    const listeners = this.socketListeners[type].get(socket);
+    listeners?.delete(listener);
+  }
+
+  emitSocket(type: 'abortRequests', socket: ClientSocket, options: WebSocketRequestAbortOptions<Schema> = {}) {
+    for (const listener of this.socketListeners[type].get(socket) ?? []) {
+      listener(options);
     }
   }
 
-  abortSocketMessages(sockets: Collection<ClientSocket> = this.sockets) {
-    const abortError = new WebSocketMessageAbortError();
+  offAny() {
+    this.channelListeners = {};
 
-    for (const socket of sockets) {
-      const listeners = this.socketListeners.messageAbort.get(socket) ?? [];
-      for (const listener of listeners) {
-        listener(abortError);
+    for (const listenersBySocket of Object.values(this.socketListeners)) {
+      for (const listeners of listenersBySocket.values()) {
+        listeners.clear();
       }
     }
   }

package/src/webSocket/WebSocketServer.ts

@@ -79,12 +79,11 @@ class WebSocketServer<Schema extends WebSocketSchema> extends WebSocketHandler<S
       return;
     }
 
-    super.removeAllChannelListeners();
-    super.abortSocketMessages();
+    super.offAny();
     await super.closeClientSockets();
 
     await closeServerSocket(this.webSocketServer, { timeout: this.socketTimeout });
-    this.webSocketServer.removeAllListeners();
+
     this.webSocketServer = undefined;
   }
 }

package/src/webSocket/errors/InvalidWebSocketMessageError.ts

@@ -0,0 +1,8 @@
+class InvalidWebSocketMessageError extends Error {
+  constructor(message: unknown) {
+    super(`Web socket message is invalid and could not be parsed: ${message}`);
+    this.name = 'InvalidWebSocketMessageError';
+  }
+}
+
+export default InvalidWebSocketMessageError;

package/src/webSocket/types.ts

@@ -1,8 +1,8 @@
-import { JSONSerialized, JSONValue } from '@zimic/http';
+import { JSONValue } from '@zimic/http';
 import { PossiblePromise } from '@zimic/utils/types';
 import type { WebSocket as ClientSocket } from 'isomorphic-ws';
 
-export interface WebSocketChannelData<Channel extends string> {
+export interface WebSocketChannelFrame<Channel extends string> {
   id: string;
   channel: Channel;
 }
@@ -10,14 +10,14 @@ export interface WebSocketChannelData<Channel extends string> {
 export interface WebSocketEventMessage<
   Schema extends WebSocketSchema,
   Channel extends WebSocketChannel<Schema> = WebSocketChannel<Schema>,
-> extends WebSocketChannelData<Channel> {
+> extends WebSocketChannelFrame<Channel> {
   data: Schema[Channel]['event'];
 }
 
 export interface WebSocketReplyMessage<
   Schema extends WebSocketSchema,
   Channel extends WebSocketChannel<Schema> = WebSocketChannel<Schema>,
-> extends WebSocketChannelData<Channel> {
+> extends WebSocketChannelFrame<Channel> {
   data: Schema[Channel]['reply'];
   requestId: string;
 }
@@ -34,16 +34,7 @@ interface BaseWebSocketSchema {
   };
 }
 
-export type WebSocketSchema<Schema extends BaseWebSocketSchema = BaseWebSocketSchema> =
-  WebSocketSchema.ConvertToStrict<Schema>;
-
-export namespace WebSocketSchema {
-  export type ConvertToStrict<Schema extends BaseWebSocketSchema> = {
-    [Channel in keyof Schema]: {
-      [Key in keyof Schema[Channel]]: JSONSerialized<Schema[Channel][Key]>;
-    };
-  };
-}
+export type WebSocketSchema<Schema extends BaseWebSocketSchema = BaseWebSocketSchema> = Schema;
 
 export type WebSocketChannel<Schema extends WebSocketSchema> = keyof Schema & string;
 
@@ -65,3 +56,7 @@ export type WebSocketReplyMessageListener<Schema extends WebSocketSchema, Channe
   message: WebSocketReplyMessage<Schema, Channel>,
   socket: ClientSocket,
 ) => PossiblePromise<void>;
+
+export type WebSocketMessageListener<Schema extends WebSocketSchema, Channel extends WebSocketChannel<Schema>> =
+  | WebSocketEventMessageListener<Schema, Channel>
+  | WebSocketReplyMessageListener<Schema, Channel>;