|
| 1 | +/** |
| 2 | + * @vitest-environment node |
| 3 | + */ |
| 4 | +import { createServer, type Server as HttpServer } from 'node:http' |
| 5 | +import type { AddressInfo } from 'node:net' |
| 6 | +import { Server } from 'socket.io' |
| 7 | +import { io as connect, type Socket } from 'socket.io-client' |
| 8 | +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' |
| 9 | +import { setupConnectionHandlers, waitForConnectionCleanup } from '@/handlers/connection' |
| 10 | +import type { AuthenticatedSocket } from '@/middleware/auth' |
| 11 | +import { MemoryRoomManager } from '@/rooms' |
| 12 | + |
| 13 | +vi.mock('@/handlers/file-doc', () => ({ cleanupFileDocForSocket: vi.fn() })) |
| 14 | +vi.mock('@/handlers/subblocks', () => ({ cleanupPendingSubblocksForSocket: vi.fn() })) |
| 15 | +vi.mock('@/handlers/variables', () => ({ cleanupPendingVariablesForSocket: vi.fn() })) |
| 16 | + |
| 17 | +describe('server shutdown connection drain', () => { |
| 18 | + let httpServer: HttpServer |
| 19 | + let io: Server |
| 20 | + let manager: MemoryRoomManager |
| 21 | + let client: Socket |
| 22 | + |
| 23 | + beforeEach(async () => { |
| 24 | + httpServer = createServer() |
| 25 | + io = new Server(httpServer, { transports: ['websocket'] }) |
| 26 | + manager = new MemoryRoomManager(io) |
| 27 | + await manager.initialize() |
| 28 | + io.on('connection', (socket) => setupConnectionHandlers(socket as AuthenticatedSocket, manager)) |
| 29 | + await new Promise<void>((resolve) => httpServer.listen(0, '127.0.0.1', resolve)) |
| 30 | + const port = (httpServer.address() as AddressInfo).port |
| 31 | + client = connect(`http://127.0.0.1:${port}`, { transports: ['websocket'], autoConnect: false }) |
| 32 | + const connected = new Promise<void>((resolve) => client.once('connect', resolve)) |
| 33 | + client.connect() |
| 34 | + await connected |
| 35 | + }) |
| 36 | + |
| 37 | + afterEach(async () => { |
| 38 | + client.disconnect() |
| 39 | + await io.close() |
| 40 | + await waitForConnectionCleanup() |
| 41 | + await manager.shutdown() |
| 42 | + vi.restoreAllMocks() |
| 43 | + }) |
| 44 | + |
| 45 | + it('keeps automatic reconnection active after transport shutdown', async () => { |
| 46 | + const disconnected = new Promise<string>((resolve) => client.once('disconnect', resolve)) |
| 47 | + await io.close() |
| 48 | + expect(await disconnected).toBe('transport close') |
| 49 | + expect(client.active).toBe(true) |
| 50 | + await waitForConnectionCleanup() |
| 51 | + }) |
| 52 | + |
| 53 | + it('waits for asynchronous presence cleanup before releasing its dependencies', async () => { |
| 54 | + let finishRemoval: (() => void) | undefined |
| 55 | + vi.spyOn(manager, 'removeSocketFromAllRooms').mockImplementation( |
| 56 | + () => |
| 57 | + new Promise((resolve) => { |
| 58 | + finishRemoval = () => resolve([]) |
| 59 | + }) |
| 60 | + ) |
| 61 | + await io.close() |
| 62 | + let drained = false |
| 63 | + const drain = waitForConnectionCleanup().then(() => { |
| 64 | + drained = true |
| 65 | + }) |
| 66 | + await Promise.resolve() |
| 67 | + expect(drained).toBe(false) |
| 68 | + expect(finishRemoval).toBeDefined() |
| 69 | + finishRemoval?.() |
| 70 | + await drain |
| 71 | + expect(drained).toBe(true) |
| 72 | + }) |
| 73 | +}) |
0 commit comments