diff --git a/backend/src/websocket/__tests__/yjsServer.test.ts b/backend/src/websocket/__tests__/yjsServer.test.ts new file mode 100644 index 0000000..0887335 --- /dev/null +++ b/backend/src/websocket/__tests__/yjsServer.test.ts @@ -0,0 +1,172 @@ +import { describe, it, expect, beforeAll, afterAll } from 'vitest' +import * as Y from 'yjs' +import { WebsocketProvider } from 'y-websocket' +import { createServer } from 'http' +import { WebSocketServer } from 'ws' +import { handleWebSocketConnection } from '../yjsServer' + +// Mock JWT verification +jest.mock('../auth', () => ({ + verifyJWT: jest.fn().mockImplementation((token) => { + if (token === 'valid-editor-token') { + return { id: 'user1', role: 'editor' } + } else if (token === 'valid-viewer-token') { + return { id: 'user2', role: 'viewer' } + } else { + throw new Error('Invalid token') + } + }) +})) + +describe('Yjs WebSocket Server', () => { + let server: any + let wss: WebSocketServer + let port: number + + beforeAll((done) => { + // Create a test HTTP server + server = createServer() + + // Create a WebSocket server + wss = new WebSocketServer({ server }) + + // Handle WebSocket connections + wss.on('connection', handleWebSocketConnection) + + // Start the server on a random port + server.listen(0, () => { + port = (server.address() as any).port + done() + }) + }) + + afterAll(() => { + // Clean up + wss.close() + server.close() + }) + + it('should sync changes between two clients', (done) => { + // Create two Yjs documents + const doc1 = new Y.Doc() + const doc2 = new Y.Doc() + + // Create WebSocket providers for both clients + const provider1 = new WebsocketProvider(`ws://localhost:${port}`, 'project:test-project', doc1) + const provider2 = new WebsocketProvider(`ws://localhost:${port}`, 'project:test-project', doc2) + + // Set up awareness for both clients + provider1.awareness.setLocalState({ user: { name: 'User 1' } }) + provider2.awareness.setLocalState({ user: { name: 'User 2' } }) + + // Wait for both clients to connect + Promise.all([ + new Promise(resolve => provider1.once('status', ({ status }: any) => status === 'connected' && resolve(true))), + new Promise(resolve => provider2.once('status', ({ status }: any) => status === 'connected' && resolve(true))) + ]).then(() => { + // Make a change in the first document + doc1.getText('content').insert(0, 'Hello World') + + // Wait for the change to sync to the second document + setTimeout(() => { + expect(doc2.getText('content').toString()).toBe('Hello World') + + // Clean up + provider1.disconnect() + provider2.disconnect() + + done() + }, 100) + }) + }) + + it('should verify JWT authentication on WebSocket connection', (done) => { + // Create a Yjs document + const doc = new Y.Doc() + + // Create a WebSocket provider with an invalid token + const provider = new WebsocketProvider(`ws://localhost:${port}?token=invalid-token`, 'project:test-project', doc) + + // Wait for the connection to be rejected + provider.once('connection-close', ({ code }: any) => { + expect(code).toBe(4000) // Invalid token error code + done() + }) + }) + + it('should prevent read-only users from sending updates', (done) => { + // Create two Yjs documents + const editorDoc = new Y.Doc() + const viewerDoc = new Y.Doc() + + // Create WebSocket providers for editor and viewer + const editorProvider = new WebsocketProvider(`ws://localhost:${port}?token=valid-editor-token`, 'project:test-project', editorDoc) + const viewerProvider = new WebsocketProvider(`ws://localhost:${port}?token=valid-viewer-token`, 'project:test-project', viewerDoc) + + // Wait for both clients to connect + Promise.all([ + new Promise(resolve => editorProvider.once('status', ({ status }: any) => status === 'connected' && resolve(true))), + new Promise(resolve => viewerProvider.once('status', ({ status }: any) => status === 'connected' && resolve(true))) + ]).then(() => { + // Try to make a change from the viewer (should be blocked) + viewerDoc.getText('content').insert(0, 'Viewer change') + + // Make a change from the editor (should be allowed) + editorDoc.getText('content').insert(0, 'Editor change') + + // Wait to see if the viewer's change gets through + setTimeout(() => { + // The editor's change should be there + expect(editorDoc.getText('content').toString()).toBe('Editor change') + + // The viewer's change should not be there + expect(viewerDoc.getText('content').toString()).toBe('') + + // Clean up + editorProvider.disconnect() + viewerProvider.disconnect() + + done() + }, 100) + }) + }) + + it('should broadcast awareness updates correctly', (done) => { + // Create two Yjs documents + const doc1 = new Y.Doc() + const doc2 = new Y.Doc() + + // Create WebSocket providers for both clients + const provider1 = new WebsocketProvider(`ws://localhost:${port}?token=valid-editor-token`, 'project:test-project', doc1) + const provider2 = new WebsocketProvider(`ws://localhost:${port}?token=valid-editor-token`, 'project:test-project', doc2) + + // Set up awareness for both clients + provider1.awareness.setLocalState({ user: { name: 'User 1', color: 'red' } }) + + // Wait for both clients to connect + Promise.all([ + new Promise(resolve => provider1.once('status', ({ status }: any) => status === 'connected' && resolve(true))), + new Promise(resolve => provider2.once('status', ({ status }: any) => status === 'connected' && resolve(true))) + ]).then(() => { + // Update awareness for client 2 + provider2.awareness.setLocalState({ user: { name: 'User 2', color: 'blue' } }) + + // Wait for awareness to sync + setTimeout(() => { + // Check that client 1 received client 2's awareness + const states = provider1.awareness.getStates() + const client2State = Array.from(states.entries()).find(([key, state]: [number, any]) => + state.user && state.user.name === 'User 2' + ) + + expect(client2State).toBeDefined() + + // Clean up + provider1.disconnect() + provider2.disconnect() + + done() + }, 100) + }) + }) +})