Merge pull request #26322 from storybookjs/norbert/ws-disconnect-event

Core: Add event when serverChannel disconnects

code/lib/channels/src/index.test.ts

@@ -1,9 +1,43 @@
 import { describe, beforeEach, it, expect, vi } from 'vitest';
 import type { ChannelTransport, Listener } from '.';
-import { Channel } from '.';
+import { Channel, WebsocketTransport } from '.';
 
 vi.useFakeTimers();
 
+const MockedWebsocket = vi.hoisted(() => {
+  const ref = { current: undefined as unknown as InstanceType<typeof MyMockedWebsocket> };
+  class MyMockedWebsocket {
+    onopen: () => void;
+
+    onmessage: (event: { data: string }) => void;
+
+    onerror: (e: any) => void;
+
+    onclose: () => void;
+
+    constructor(url: string) {
+      this.onopen = vi.fn();
+      this.onmessage = vi.fn();
+      this.onerror = vi.fn();
+      this.onclose = vi.fn();
+
+      ref.current = this;
+    }
+
+    send(data: string) {
+      this.onmessage({ data });
+    }
+  }
+  return { MyMockedWebsocket, ref };
+});
+
+vi.mock('@storybook/global', () => ({
+  global: {
+    ...global,
+    WebSocket: MockedWebsocket.MyMockedWebsocket,
+  },
+}));
+
 describe('Channel', () => {
   let transport: ChannelTransport;
   let channel: Channel;
@@ -232,3 +266,78 @@ describe('Channel', () => {
     });
   });
 });
+
+describe('WebsocketTransport', () => {
+  it('should connect', async () => {
+    const onError = vi.fn();
+    const handler = vi.fn();
+
+    const transport = new WebsocketTransport({
+      url: 'ws://localhost:6006',
+      page: 'preview',
+      onError,
+    });
+
+    transport.setHandler(handler);
+    MockedWebsocket.ref.current.onopen();
+
+    expect(handler).toHaveBeenCalledTimes(0);
+  });
+  it('should send message upon disconnect', async () => {
+    const onError = vi.fn();
+    const handler = vi.fn();
+
+    const transport = new WebsocketTransport({
+      url: 'ws://localhost:6006',
+      page: 'preview',
+      onError,
+    });
+
+    transport.setHandler(handler);
+    MockedWebsocket.ref.current.onclose();
+
+    expect(handler.mock.calls[0][0]).toMatchInlineSnapshot(`
+      {
+        "args": [],
+        "from": "preview",
+        "type": "channelWSDisconnect",
+      }
+    `);
+  });
+  it('should send message when send', async () => {
+    const onError = vi.fn();
+    const handler = vi.fn();
+
+    const transport = new WebsocketTransport({
+      url: 'ws://localhost:6006',
+      page: 'preview',
+      onError,
+    });
+
+    transport.setHandler(handler);
+    MockedWebsocket.ref.current.send('{ "type": "test", "args": [], "from": "preview" }');
+
+    expect(handler.mock.calls[0][0]).toMatchInlineSnapshot(`
+      {
+        "args": [],
+        "from": "preview",
+        "type": "test",
+      }
+    `);
+  });
+  it('should call onError handler', async () => {
+    const onError = vi.fn();
+    const handler = vi.fn();
+
+    const transport = new WebsocketTransport({
+      url: 'ws://localhost:6006',
+      page: 'preview',
+      onError,
+    });
+
+    transport.setHandler(handler);
+    MockedWebsocket.ref.current.onerror(new Error('testError'));
+
+    expect(onError.mock.calls[0][0]).toMatchInlineSnapshot(`[Error: testError]`);
+  });
+});

code/lib/channels/src/index.ts

@@ -35,7 +35,7 @@ export function createBrowserChannel({ page, extraTransports = [] }: Options): C
     const { hostname, port } = window.location;
     const channelUrl = `${protocol}://${hostname}:${port}/storybook-server-channel`;
 
-    transports.push(new WebsocketTransport({ url: channelUrl, onError: () => {} }));
+    transports.push(new WebsocketTransport({ url: channelUrl, onError: () => {}, page }));
   }
 
   return new Channel({ transports });

code/lib/channels/src/websocket/index.ts

@@ -5,13 +5,14 @@ import { global } from '@storybook/global';
 import { isJSON, parse, stringify } from 'telejson';
 import invariant from 'tiny-invariant';
 
-import type { ChannelTransport, ChannelHandler } from '../types';
+import * as EVENTS from '@storybook/core-events';
+import type { ChannelTransport, ChannelHandler, Config } from '../types';
 
 const { WebSocket } = global;
 
 type OnError = (message: Event) => void;
 
-interface WebsocketTransportArgs {
+interface WebsocketTransportArgs extends Partial<Config> {
   url: string;
   onError: OnError;
 }
@@ -25,7 +26,7 @@ export class WebsocketTransport implements ChannelTransport {
 
   private isReady = false;
 
-  constructor({ url, onError }: WebsocketTransportArgs) {
+  constructor({ url, onError, page }: WebsocketTransportArgs) {
     this.socket = new WebSocket(url);
     this.socket.onopen = () => {
       this.isReady = true;
@@ -41,6 +42,10 @@ export class WebsocketTransport implements ChannelTransport {
         onError(e);
       }
     };
+    this.socket.onclose = () => {
+      invariant(this.handler, 'WebsocketTransport handler should be set');
+      this.handler({ type: EVENTS.CHANNEL_WS_DISCONNECT, args: [], from: page || 'preview' });
+    };
   }
 
   setHandler(handler: ChannelHandler) {

code/lib/core-events/src/index.ts

@@ -1,5 +1,6 @@
 // eslint-disable-next-line @typescript-eslint/naming-convention
 enum events {
+  CHANNEL_WS_DISCONNECT = 'channelWSDisconnect',
   CHANNEL_CREATED = 'channelCreated',
   // There was an error executing the config, likely an bug in the user's preview.js
   CONFIG_ERROR = 'configError',
@@ -80,6 +81,7 @@ export default events;
 // Enables: `import * as Events from ...` or `import { CHANNEL_CREATED } as Events from ...`
 // This is the preferred method
 export const {
+  CHANNEL_WS_DISCONNECT,
   CHANNEL_CREATED,
   CONFIG_ERROR,
   CURRENT_STORY_WAS_SET,

code/ui/manager/src/globals/exports.ts

@@ -130,6 +130,7 @@ export default {
   ],
   '@storybook/core-events': [
     'CHANNEL_CREATED',
+    'CHANNEL_WS_DISCONNECT',
     'CONFIG_ERROR',
     'CURRENT_STORY_WAS_SET',
     'DOCS_PREPARED',