fix: add websocket logging

Author: kettanaito

src/browser/setup-worker.ts

@@ -18,8 +18,7 @@ import { InterceptorSource } from '~/core/new/sources/interceptor-source'
 
 const DEFAULT_WORKER_URL = '/mockServiceWorker.js'
 
-export interface SetupWorkerApi
-  extends NetworkHandlersApi<HttpHandler | WebSocketHandler> {
+export interface SetupWorkerApi extends NetworkHandlersApi {
   start: (options?: SetupWorkerStartOptions) => Promise<void>
   stop: () => void
 }
@@ -48,7 +47,7 @@ export type FindWorker = (
 export function setupWorker(
   ...handlers: Array<HttpHandler | WebSocketHandler>
 ): SetupWorkerApi {
-  let network: NetworkApi<HttpHandler | WebSocketHandler>
+  let network: NetworkApi
 
   return {
     async start(options) {

src/core/handlers/WebSocketHandler.ts

@@ -13,6 +13,7 @@ import {
 } from '../utils/matching/matchRequestUrl'
 import { getCallFrame } from '../utils/internal/getCallFrame'
 import type { HandlerKind } from './common'
+import { attachWebSocketLogger } from '../ws/utils/attachWebSocketLogger'
 
 type WebSocketHandlerParsedResult = {
   match: Match
@@ -138,6 +139,10 @@ export class WebSocketHandler {
     // This is what the developer adds listeners for.
     return this[kEmitter].emit('connection', connection)
   }
+
+  public log(connection: WebSocketConnectionData): void {
+    attachWebSocketLogger(connection)
+  }
 }
 
 function createStopPropagationListener(handler: WebSocketHandler) {

src/core/new/define-network.ts

@@ -14,7 +14,7 @@ import {
 } from './on-unhandled-frame'
 import { pipeEvents } from '../utils/internal/pipeEvents'
 
-export interface DefineNetworkOptions<Handler extends AnyHandler> {
+export interface DefineNetworkOptions {
   /**
    * A list of network sources.
    */
@@ -23,7 +23,7 @@ export interface DefineNetworkOptions<Handler extends AnyHandler> {
   /**
    * A list of the initial handlers.
    */
-  handlers: Array<Handler>
+  handlers: Array<AnyHandler>
 
   /**
    * A custom handlers controller to use.
@@ -39,8 +39,7 @@ export interface DefineNetworkOptions<Handler extends AnyHandler> {
   quiet?: boolean
 }
 
-export interface NetworkApi<Handler extends AnyHandler>
-  extends NetworkHandlersApi<Handler> {
+export interface NetworkApi extends NetworkHandlersApi {
   /**
    * Enables the network interception.
    */
@@ -57,16 +56,14 @@ export interface NetworkApi<Handler extends AnyHandler>
   events: Emitter<any>
 }
 
-export interface NetworkHandlersApi<Handler extends AnyHandler> {
-  use(...handlers: Array<Handler>): void
-  resetHandlers(...nextHandlers: Array<Handler>): void
+export interface NetworkHandlersApi {
+  use(...handlers: Array<AnyHandler>): void
+  resetHandlers(...nextHandlers: Array<AnyHandler>): void
   restoreHandlers(): void
-  listHandlers(): ReadonlyArray<Handler>
+  listHandlers(): ReadonlyArray<AnyHandler>
 }
 
-export function defineNetwork<Handler extends AnyHandler>(
-  options: DefineNetworkOptions<Handler>,
-): NetworkApi<Handler> {
+export function defineNetwork(options: DefineNetworkOptions): NetworkApi {
   const source = NetworkSource.from(...options.sources)
   const events = new Emitter()
 

src/core/new/resolve-network-frame.ts

@@ -126,14 +126,9 @@ async function resolveHttpNetworkFrame(
   frame.events.emit('request:end', { request, requestId })
 
   if (!flow.quiet) {
-    /**
-     * @fixme This doesn't belong here. Different network APIs might choose
-     * to handle logging differently (e.g. `setupServer` doesn't log at all).
-     * This likely belongs in an abstract method on `defineNetwork()` or something.
-     */
     // Log mocked responses. Use the Network tab to observe the original network.
     handler.log({
-      request: requestCloneForLogs,
+      request: requestCloneForLogs!,
       response,
       parsedResult,
     })
@@ -148,7 +143,7 @@ async function resolveHttpNetworkFrame(
 async function resolveWebSocketNetworkFrame(
   frame: WebSocketNetworkFrame,
   handlers: Array<AnyHandler>,
-  _flow: ResolveNetworkFlow,
+  flow: ResolveNetworkFlow,
 ): Promise<void> {
   const { connection } = frame.data
   const eventHandlers = handlers.filter(isHandlerKind('EventHandler'))
@@ -159,21 +154,32 @@ async function resolveWebSocketNetworkFrame(
   })
 
   if (eventHandlers.length > 0) {
-    await Promise.all(
-      eventHandlers.map((handler) => {
+    const matches = await Promise.all<boolean>(
+      eventHandlers.map(async (handler) => {
         // Foward the connection data to every WebSocket handler.
         // This is equivalent to dispatching the connection event
         // onto multiple listeners.
-        return handler.run(connection)
+        return handler.run(connection).then((matches) => {
+          if (matches) {
+            flow?.handled?.({ frame, handler })
+
+            if (!flow?.quiet) {
+              handler.log(connection)
+            }
+          }
+
+          return matches
+        })
       }),
     )
 
+    if (matches.every((match) => !match)) {
+      flow?.unhandled?.()
+    }
+
     return
   }
 
-  /**
-   * @todo Support WebSocket logging, somehow.
-   */
-
   frame.passthrough()
+  flow.unhandled?.()
 }