directus

1
import type { Accountability } from '@directus/types';
2
import { afterEach, describe, expect, test, vi } from 'vitest';
3
import { WebSocketController, getWebSocketController } from '../websocket/controllers/index.js';
4
import type { WebSocketClient } from '../websocket/types.js';
5
import { WebSocketService } from './websocket.js';
6

7
vi.mock('../emitter');
8
vi.mock('../websocket/controllers/index');
9

10
// This is required because logger uses global env which is imported before the tests run. Can be
11
// reduce to just mock the file when logger is also using useLogger everywhere @TODO
12
vi.mock('@directus/env', () => ({
13
	useEnv: vi.fn().mockReturnValue({
14
		WEBSOCKETS_ENABLED: true,
15
		WEBSOCKETS_REST_ENABLED: true,
16
		EXTENSIONS_PATH: './extensions',
17
		EMAIL_TEMPLATES_PATH: './templates',
18
	}),
19
}));
20

21
afterEach(() => {
22
	vi.clearAllMocks();
23
});
24

25
function mockClient(accountability: Accountability | null = null) {
26
	return {
27
		on: vi.fn(),
28
		off: vi.fn(),
29
		send: vi.fn(),
30
		close: vi.fn(),
31
		accountability,
32
	} as unknown as WebSocketClient;
33
}
34

35
describe('WebSocketService', () => {
36
	test('get clients', () => {
37
		vi.mocked(getWebSocketController).mockReturnValue({
38
			clients: new Set([mockClient(), mockClient(), mockClient()]),
39
		} as unknown as WebSocketController);
40

41
		const wsService = new WebSocketService();
42
		expect(wsService.clients().size).toBe(3);
43
	});
44

45
	test('broadcast', () => {
46
		const clients = new Set([mockClient(), mockClient(), mockClient()]);
47
		const message = 'test 123';
48

49
		vi.mocked(getWebSocketController).mockReturnValue({ clients } as unknown as WebSocketController);
50

51
		const wsService = new WebSocketService();
52
		wsService.broadcast(message);
53

54
		for (const client of clients) {
55
			expect(client.send).toBeCalledWith(message);
56
		}
57
	});
58

59
	test('broadcast with role filter', () => {
60
		const clients = [mockClient({ user: 'test', role: 'test' }), mockClient({ user: 'test2', role: 'test2' })];
61
		const message = 'test 123';
62

63
		vi.mocked(getWebSocketController).mockReturnValue({ clients: new Set(clients) } as unknown as WebSocketController);
64

65
		const wsService = new WebSocketService();
66
		wsService.broadcast(message, { role: 'test' });
67

68
		expect(clients[0]!.send).toBeCalledWith(message);
69
		expect(clients[1]!.send).not.toBeCalled();
70
	});
71

72
	test('broadcast with user filter', () => {
73
		const clients = [mockClient({ user: 'test', role: 'test' }), mockClient({ user: 'test2', role: 'test2' })];
74
		const message = 'test 123';
75

76
		vi.mocked(getWebSocketController).mockReturnValue({ clients: new Set(clients) } as unknown as WebSocketController);
77

78
		const wsService = new WebSocketService();
79
		wsService.broadcast(message, { user: 'test2' });
80

81
		expect(clients[0]!.send).not.toBeCalled();
82
		expect(clients[1]!.send).toBeCalledWith(message);
83
	});
84
});
85