feat: Add tests for Yjs WebSocket server

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)
+    })
+  })
+})