From 788aaac74729a17f9d08eeba02f9d8905a3451a6 Mon Sep 17 00:00:00 2001 From: Norbert de Langen Date: Tue, 5 Mar 2024 13:34:46 +0100 Subject: [PATCH 1/2] add functionality of emitting an event & and test for websocket-transport --- code/lib/channels/src/index.test.ts | 125 ++++++++++++++++++++++- code/lib/channels/src/index.ts | 2 +- code/lib/channels/src/websocket/index.ts | 14 ++- code/lib/core-events/src/index.ts | 2 + code/ui/manager/src/globals/exports.ts | 1 + 5 files changed, 139 insertions(+), 5 deletions(-) diff --git a/code/lib/channels/src/index.test.ts b/code/lib/channels/src/index.test.ts index 07f6d605af10..04c295d98cff 100644 --- a/code/lib/channels/src/index.test.ts +++ b/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 }; + 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,92 @@ 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).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).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).toMatchInlineSnapshot(` + [ + [ + [Error: testError], + ], + ] + `); + }); +}); diff --git a/code/lib/channels/src/index.ts b/code/lib/channels/src/index.ts index 7942d57fa4f3..80a865f31904 100644 --- a/code/lib/channels/src/index.ts +++ b/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 }); diff --git a/code/lib/channels/src/websocket/index.ts b/code/lib/channels/src/websocket/index.ts index a46df1c28610..790f69a8deaa 100644 --- a/code/lib/channels/src/websocket/index.ts +++ b/code/lib/channels/src/websocket/index.ts @@ -5,13 +5,16 @@ import { global } from '@storybook/global'; import { isJSON, parse, stringify } from 'telejson'; import invariant from 'tiny-invariant'; -import type { ChannelTransport, ChannelHandler } from '../types'; +// I tried to use an import statement, but it didn't work +const { CHANNEL_WS_DISCONNECT } = require('@storybook/core-events'); + +import type { ChannelTransport, ChannelHandler, Config } from '../types'; const { WebSocket } = global; type OnError = (message: Event) => void; -interface WebsocketTransportArgs { +interface WebsocketTransportArgs extends Partial { url: string; onError: OnError; } @@ -25,7 +28,8 @@ export class WebsocketTransport implements ChannelTransport { private isReady = false; - constructor({ url, onError }: WebsocketTransportArgs) { + constructor({ url, onError, page }: WebsocketTransportArgs) { + console.log({ WebSocket }); this.socket = new WebSocket(url); this.socket.onopen = () => { this.isReady = true; @@ -41,6 +45,10 @@ export class WebsocketTransport implements ChannelTransport { onError(e); } }; + this.socket.onclose = () => { + invariant(this.handler, 'WebsocketTransport handler should be set'); + this.handler({ type: CHANNEL_WS_DISCONNECT, args: [], from: page || 'preview' }); + }; } setHandler(handler: ChannelHandler) { diff --git a/code/lib/core-events/src/index.ts b/code/lib/core-events/src/index.ts index 6d98e6291200..7542d0aa4e57 100644 --- a/code/lib/core-events/src/index.ts +++ b/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, diff --git a/code/ui/manager/src/globals/exports.ts b/code/ui/manager/src/globals/exports.ts index b8e2d5b14868..079340369fcb 100644 --- a/code/ui/manager/src/globals/exports.ts +++ b/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', From ff54701ed7e9eadd1b8ceda4e6c042285c59c9e9 Mon Sep 17 00:00:00 2001 From: Norbert de Langen Date: Tue, 5 Mar 2024 14:11:05 +0100 Subject: [PATCH 2/2] fix --- code/lib/channels/src/index.test.ts | 40 ++++++++---------------- code/lib/channels/src/websocket/index.ts | 7 ++--- 2 files changed, 15 insertions(+), 32 deletions(-) diff --git a/code/lib/channels/src/index.test.ts b/code/lib/channels/src/index.test.ts index 04c295d98cff..f99e04f6099c 100644 --- a/code/lib/channels/src/index.test.ts +++ b/code/lib/channels/src/index.test.ts @@ -296,16 +296,12 @@ describe('WebsocketTransport', () => { transport.setHandler(handler); MockedWebsocket.ref.current.onclose(); - expect(handler.mock.calls).toMatchInlineSnapshot(` - [ - [ - { - "args": [], - "from": "preview", - "type": "channelWSDisconnect", - }, - ], - ] + expect(handler.mock.calls[0][0]).toMatchInlineSnapshot(` + { + "args": [], + "from": "preview", + "type": "channelWSDisconnect", + } `); }); it('should send message when send', async () => { @@ -321,16 +317,12 @@ describe('WebsocketTransport', () => { transport.setHandler(handler); MockedWebsocket.ref.current.send('{ "type": "test", "args": [], "from": "preview" }'); - expect(handler.mock.calls).toMatchInlineSnapshot(` - [ - [ - { - "args": [], - "from": "preview", - "type": "test", - }, - ], - ] + expect(handler.mock.calls[0][0]).toMatchInlineSnapshot(` + { + "args": [], + "from": "preview", + "type": "test", + } `); }); it('should call onError handler', async () => { @@ -346,12 +338,6 @@ describe('WebsocketTransport', () => { transport.setHandler(handler); MockedWebsocket.ref.current.onerror(new Error('testError')); - expect(onError.mock.calls).toMatchInlineSnapshot(` - [ - [ - [Error: testError], - ], - ] - `); + expect(onError.mock.calls[0][0]).toMatchInlineSnapshot(`[Error: testError]`); }); }); diff --git a/code/lib/channels/src/websocket/index.ts b/code/lib/channels/src/websocket/index.ts index 790f69a8deaa..0cc73345507d 100644 --- a/code/lib/channels/src/websocket/index.ts +++ b/code/lib/channels/src/websocket/index.ts @@ -5,9 +5,7 @@ import { global } from '@storybook/global'; import { isJSON, parse, stringify } from 'telejson'; import invariant from 'tiny-invariant'; -// I tried to use an import statement, but it didn't work -const { CHANNEL_WS_DISCONNECT } = require('@storybook/core-events'); - +import * as EVENTS from '@storybook/core-events'; import type { ChannelTransport, ChannelHandler, Config } from '../types'; const { WebSocket } = global; @@ -29,7 +27,6 @@ export class WebsocketTransport implements ChannelTransport { private isReady = false; constructor({ url, onError, page }: WebsocketTransportArgs) { - console.log({ WebSocket }); this.socket = new WebSocket(url); this.socket.onopen = () => { this.isReady = true; @@ -47,7 +44,7 @@ export class WebsocketTransport implements ChannelTransport { }; this.socket.onclose = () => { invariant(this.handler, 'WebsocketTransport handler should be set'); - this.handler({ type: CHANNEL_WS_DISCONNECT, args: [], from: page || 'preview' }); + this.handler({ type: EVENTS.CHANNEL_WS_DISCONNECT, args: [], from: page || 'preview' }); }; }