Skip to content

Commit

Permalink
Merge pull request #26322 from storybookjs/norbert/ws-disconnect-event
Browse files Browse the repository at this point in the history
Core: Add event when serverChannel disconnects
  • Loading branch information
shilman authored Mar 5, 2024
2 parents 6c71508 + b709d27 commit 85c4d9c
Show file tree
Hide file tree
Showing 5 changed files with 122 additions and 5 deletions.
111 changes: 110 additions & 1 deletion code/lib/channels/src/index.test.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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]`);
});
});
2 changes: 1 addition & 1 deletion code/lib/channels/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 });
Expand Down
11 changes: 8 additions & 3 deletions code/lib/channels/src/websocket/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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;
Expand All @@ -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) {
Expand Down
2 changes: 2 additions & 0 deletions code/lib/core-events/src/index.ts
Original file line number Diff line number Diff line change
@@ -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',
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions code/ui/manager/src/globals/exports.ts
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ export default {
],
'@storybook/core-events': [
'CHANNEL_CREATED',
'CHANNEL_WS_DISCONNECT',
'CONFIG_ERROR',
'CURRENT_STORY_WAS_SET',
'DOCS_PREPARED',
Expand Down

0 comments on commit 85c4d9c

Please sign in to comment.