@dotdo/postgres 0.1.0

package/src/worker/cdc-ws.test.ts

@@ -0,0 +1,1009 @@
+/**
+ * RED PHASE: Flaky CDC WebSocket Transport Tests
+ * Task: postgres-dnid - Fix flaky CDC WebSocket transport tests
+ *
+ * These tests are designed to EXPOSE timing-related flakiness in the CDC WebSocket transport.
+ * They test:
+ * 1. Timing-related race conditions
+ * 2. Reconnection logic under various conditions
+ * 3. Message ordering guarantees
+ * 4. Connection state management
+ *
+ * Expected: These tests should initially FAIL to identify the flaky behavior.
+ */
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
+import {
+  CDCWebSocketHandler,
+  createCDCWebSocketHandler,
+} from './cdc-websocket'
+import { CDCManager } from './cdc'
+import type { WSMessage, CDCWebSocketConfig } from './types'
+
+// ============================================================================
+// Mock Infrastructure
+// ============================================================================
+
+interface MockWebSocket {
+  id: string
+  send: ReturnType<typeof vi.fn>
+  close: ReturnType<typeof vi.fn>
+  readyState: number
+  serializeAttachment: ReturnType<typeof vi.fn>
+  deserializeAttachment: ReturnType<typeof vi.fn>
+  // Simulate WebSocket state transitions
+  simulateClose: (code: number, reason: string) => void
+  simulateError: (error: Error) => void
+  setReadyState: (state: number) => void
+}
+
+const WS_CONNECTING = 0
+const WS_OPEN = 1
+const WS_CLOSING = 2
+const WS_CLOSED = 3
+
+let mockWebSocketIdCounter = 0
+
+const createMockWebSocket = (id?: string): MockWebSocket => {
+  const wsId = id ?? `ws-${++mockWebSocketIdCounter}`
+  let readyState = WS_OPEN
+
+  const ws: MockWebSocket = {
+    id: wsId,
+    send: vi.fn(),
+    close: vi.fn(),
+    get readyState() {
+      return readyState
+    },
+    serializeAttachment: vi.fn(),
+    deserializeAttachment: vi.fn(() => ({ subscriptions: [] })),
+    simulateClose: (code: number, reason: string) => {
+      readyState = WS_CLOSED
+      ws.close(code, reason)
+    },
+    simulateError: (_error: Error) => {
+      readyState = WS_CLOSING
+    },
+    setReadyState: (state: number) => {
+      readyState = state
+    },
+  }
+  return ws
+}
+
+const createMockState = () => {
+  const webSockets: MockWebSocket[] = []
+  return {
+    id: { toString: () => 'test-do-id' },
+    acceptWebSocket: vi.fn((ws: MockWebSocket) => {
+      webSockets.push(ws)
+      return ws
+    }),
+    getWebSockets: vi.fn(() => webSockets),
+    storage: {
+      get: vi.fn(),
+      put: vi.fn(),
+    },
+  }
+}
+
+const createMockExecutor = () => ({
+  executeQuery: vi.fn(async () => ({
+    rows: [{ id: 1 }],
+    fields: [{ name: 'id', dataTypeID: 23 }],
+    rowCount: 1,
+    durationMs: 1.5,
+  })),
+})
+
+const createMockBroadcaster = () => ({
+  broadcast: vi.fn(),
+})
+
+// Helper to parse sent messages
+const getSentMessages = (ws: MockWebSocket): WSMessage[] => {
+  return ws.send.mock.calls.map((call: [string]) => JSON.parse(call[0]))
+}
+
+// Helper to get last sent message
+const getLastSentMessage = (ws: MockWebSocket): WSMessage | undefined => {
+  const messages = getSentMessages(ws)
+  return messages[messages.length - 1]
+}
+
+// Helper to wait for async operations (works with both real and fake timers)
+const flushPromises = async () => {
+  // Use queueMicrotask for microtask-based flushing that works with fake timers
+  await new Promise(resolve => queueMicrotask(resolve))
+  // Also process any pending promise jobs
+  await Promise.resolve()
+}
+
+// Helper to advance timers and flush (compatible with vi.useFakeTimers())
+// Uses advanceTimersByTimeAsync to only advance by the specified amount,
+// not run ALL timers (which could cause infinite loops with setInterval)
+const advanceTimersAndFlush = async (ms: number) => {
+  // Advance fake timers by the specified amount (runs timers that become due)
+  await vi.advanceTimersByTimeAsync(ms)
+  // Flush any remaining microtasks
+  await flushPromises()
+}
+
+/**
+ * Helper to handle WebSocket upgrade in Node.js test environment.
+ * In Cloudflare Workers, Response status 101 is valid.
+ * In Node.js, we need to catch the error and verify acceptWebSocket was called.
+ */
+const safeHandleUpgrade = (
+  handler: CDCWebSocketHandler,
+  mockState: ReturnType<typeof createMockState>,
+  ws: MockWebSocket
+): void => {
+  const mockPair = { 0: createMockWebSocket(), 1: ws }
+  const request = new Request('http://localhost/cdc', {
+    headers: { Upgrade: 'websocket' },
+  })
+
+  try {
+    const response = handler.handleUpgrade(request, mockPair as any)
+    // In real Cloudflare Workers, this would return 101
+    if (response.status === 101) {
+      return
+    }
+  } catch (e) {
+    // Expected in Node.js - status 101 not allowed in fetch Response
+    // Verify acceptWebSocket was called
+    if (mockState.acceptWebSocket.mock.calls.length === 0) {
+      throw e
+    }
+  }
+}
+
+// ============================================================================
+// Test Suites
+// ============================================================================
+
+describe('CDC WebSocket Transport - Flaky Tests', () => {
+  let handler: CDCWebSocketHandler
+  let mockState: ReturnType<typeof createMockState>
+  let mockExecutor: ReturnType<typeof createMockExecutor>
+  let timers: ReturnType<typeof vi.useFakeTimers> | null = null
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockWebSocketIdCounter = 0
+    mockState = createMockState()
+    mockExecutor = createMockExecutor()
+  })
+
+  afterEach(async () => {
+    // CRITICAL: Clean up timers to prevent leaks
+    if (timers) {
+      vi.useRealTimers()
+      timers = null
+    }
+    // Allow any pending microtasks to complete
+    await flushPromises()
+  })
+
+  // ==========================================================================
+  // 1. Timing-Related Race Conditions
+  // ==========================================================================
+
+  describe('Timing-Related Race Conditions', () => {
+    it('should handle heartbeat timeout during active message processing', async () => {
+      /**
+       * FLAKY SCENARIO: Heartbeat timer fires while a long-running query is being processed.
+       * The heartbeat check might incorrectly close the connection.
+       */
+      timers = vi.useFakeTimers()
+
+      const config: CDCWebSocketConfig = {
+        heartbeatInterval: 100,
+        heartbeatTimeout: 50,
+      }
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, undefined, config)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Simulate slow query that takes longer than heartbeat timeout
+      mockExecutor.executeQuery.mockImplementation(async () => {
+        // Simulate query taking 200ms (longer than heartbeat interval)
+        await new Promise(resolve => setTimeout(resolve, 200))
+        return { rows: [], fields: [], rowCount: 0, durationMs: 200 }
+      })
+
+      // Start a query
+      const queryPromise = handler.handleMessage(ws as any, JSON.stringify({
+        type: 'query',
+        id: 'slow-query-1',
+        payload: { sql: 'SELECT pg_sleep(1)' },
+      }))
+
+      // Advance past heartbeat interval + timeout
+      await advanceTimersAndFlush(160)
+
+      // Complete the query
+      await advanceTimersAndFlush(50)
+      await queryPromise
+
+      // Connection should still be active after slow query completes
+      expect(ws.readyState).toBe(WS_OPEN)
+      expect(ws.close).not.toHaveBeenCalled()
+    })
+
+    it('should not race between heartbeat timer and explicit close', async () => {
+      /**
+       * FLAKY SCENARIO: Client closes connection exactly when heartbeat timer fires.
+       * Could cause double-close or accessing closed connection state.
+       */
+      timers = vi.useFakeTimers()
+
+      const config: CDCWebSocketConfig = {
+        heartbeatInterval: 100,
+        heartbeatTimeout: 50,
+      }
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, undefined, config)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Advance to just before heartbeat
+      await advanceTimersAndFlush(99)
+
+      // Explicitly close the connection
+      ws.simulateClose(1000, 'Client closing')
+      handler.handleClose(ws as any, 1000, 'Client closing')
+
+      // Now advance past heartbeat - timer should not cause errors
+      await advanceTimersAndFlush(10)
+
+      // Should not throw or cause issues
+      expect(handler.getConnectedCount()).toBe(0)
+    })
+
+    it('should handle rapid subscribe/unsubscribe cycles', async () => {
+      /**
+       * FLAKY SCENARIO: Rapid subscription changes can cause race conditions
+       * in subscription tracking.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+      ws.send.mockClear()
+
+      // Rapid subscribe/unsubscribe cycles
+      const operations: Promise<void>[] = []
+      const subscriptionIds: string[] = []
+
+      for (let i = 0; i < 10; i++) {
+        // Subscribe
+        operations.push(handler.handleMessage(ws as any, JSON.stringify({
+          type: 'subscribe',
+          id: `sub-${i}`,
+          payload: { table: 'users' },
+        })))
+      }
+
+      await Promise.all(operations)
+
+      // Collect subscription IDs from responses
+      const messages = getSentMessages(ws)
+      for (const msg of messages) {
+        if (msg.type === 'subscribed' && msg.payload) {
+          const payload = msg.payload as { subscriptionId: string }
+          subscriptionIds.push(payload.subscriptionId)
+        }
+      }
+
+      // Now unsubscribe all rapidly
+      ws.send.mockClear()
+      const unsubOps: Promise<void>[] = []
+      for (const subId of subscriptionIds) {
+        unsubOps.push(handler.handleMessage(ws as any, JSON.stringify({
+          type: 'unsubscribe',
+          id: `unsub-${subId}`,
+          payload: { subscriptionId: subId },
+        })))
+      }
+
+      await Promise.all(unsubOps)
+
+      // All subscriptions should be removed
+      expect(handler.getSubscriptionCount()).toBe(0)
+    })
+
+    it('should handle messages arriving during connection close', async () => {
+      /**
+       * FLAKY SCENARIO: Messages arrive after close has been initiated but
+       * before state is fully cleaned up.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Subscribe to a table
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      // Start closing
+      ws.setReadyState(WS_CLOSING)
+
+      // Try to send another message during close
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-2',
+        payload: { table: 'orders' },
+      }))
+
+      // Complete close
+      ws.simulateClose(1000, 'Normal closure')
+      handler.handleClose(ws as any, 1000, 'Normal closure')
+
+      // Should not throw and state should be clean
+      expect(handler.getConnectedCount()).toBe(0)
+    })
+  })
+
+  // ==========================================================================
+  // 2. Reconnection Logic
+  // ==========================================================================
+
+  describe('Reconnection Logic', () => {
+    it('should restore subscriptions on reconnect', async () => {
+      /**
+       * FLAKY SCENARIO: Subscription restoration during reconnect might miss
+       * subscriptions if they are being modified concurrently.
+       */
+      timers = vi.useFakeTimers()
+
+      const config: CDCWebSocketConfig = {
+        reconnectDelay: 100,
+        maxReconnectAttempts: 3,
+      }
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, undefined, config)
+
+      // First connection
+      const ws1 = createMockWebSocket('ws-1')
+      safeHandleUpgrade(handler, mockState, ws1)
+
+      // Get connection ID from welcome message
+      const welcome1 = getLastSentMessage(ws1) as any
+      expect(welcome1?.type).toBe('welcome')
+      const connectionId = welcome1?.payload?.connectionId
+
+      // Subscribe to tables
+      await handler.handleMessage(ws1 as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      const subResponse = getLastSentMessage(ws1) as any
+      const subscriptionId = subResponse?.payload?.subscriptionId
+
+      // Close first connection
+      ws1.simulateClose(1006, 'Abnormal closure')
+      handler.handleClose(ws1 as any, 1006, 'Abnormal closure')
+
+      // Second connection - reconnect
+      const ws2 = createMockWebSocket('ws-2')
+      safeHandleUpgrade(handler, mockState, ws2)
+      ws2.send.mockClear()
+
+      // Attempt reconnection with previous connection ID
+      await handler.handleMessage(ws2 as any, JSON.stringify({
+        type: 'reconnect',
+        id: 'reconnect-1',
+        payload: {
+          connectionId: connectionId,
+          subscriptionIds: [subscriptionId],
+        },
+      }))
+
+      // Should receive welcome with restored subscriptions
+      const messages = getSentMessages(ws2)
+      const welcomeMsg = messages.find(m => m.type === 'welcome')
+      expect(welcomeMsg).toBeDefined()
+    })
+
+    it('should handle reconnect with stale connection ID', async () => {
+      /**
+       * FLAKY SCENARIO: Reconnect request uses a connection ID that has
+       * already been cleaned up.
+       */
+      timers = vi.useFakeTimers()
+
+      const config: CDCWebSocketConfig = {
+        reconnectDelay: 100,
+        maxReconnectAttempts: 3,
+      }
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, undefined, config)
+
+      // First connection
+      const ws1 = createMockWebSocket('ws-1')
+      safeHandleUpgrade(handler, mockState, ws1)
+
+      const welcome1 = getLastSentMessage(ws1) as any
+      const oldConnectionId = welcome1?.payload?.connectionId
+
+      // Close first connection
+      ws1.simulateClose(1006, 'Abnormal closure')
+      handler.handleClose(ws1 as any, 1006, 'Abnormal closure')
+
+      // Wait for cleanup timeout to expire
+      await advanceTimersAndFlush(config.reconnectDelay! * config.maxReconnectAttempts! + 100)
+
+      // New connection - reconnect with stale ID
+      const ws2 = createMockWebSocket('ws-2')
+      safeHandleUpgrade(handler, mockState, ws2)
+      ws2.send.mockClear()
+
+      // Attempt reconnection with stale connection ID
+      await handler.handleMessage(ws2 as any, JSON.stringify({
+        type: 'reconnect',
+        id: 'reconnect-1',
+        payload: {
+          connectionId: oldConnectionId,
+          subscriptionIds: ['stale-sub-id'],
+        },
+      }))
+
+      // Should still work (graceful degradation) - just won't restore subscriptions
+      const messages = getSentMessages(ws2)
+      expect(messages.some(m => m.type === 'welcome' || m.type === 'error')).toBe(true)
+    })
+
+    it('should buffer changes during reconnection window', async () => {
+      /**
+       * FLAKY SCENARIO: Changes that occur during the reconnection window
+       * should be buffered and delivered on reconnect.
+       */
+      timers = vi.useFakeTimers()
+
+      const config: CDCWebSocketConfig = {
+        reconnectDelay: 100,
+        maxReconnectAttempts: 3,
+        bufferChanges: true,
+        maxBufferedChanges: 100,
+      }
+      const broadcaster = createMockBroadcaster()
+      const cdcManager = new CDCManager(broadcaster as any)
+
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, cdcManager, config)
+
+      // First connection
+      const ws1 = createMockWebSocket('ws-1')
+      safeHandleUpgrade(handler, mockState, ws1)
+
+      const welcome1 = getLastSentMessage(ws1) as any
+      const connectionId = welcome1?.payload?.connectionId
+
+      // Subscribe
+      await handler.handleMessage(ws1 as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      const subResponse = getLastSentMessage(ws1) as any
+      const subscriptionId = subResponse?.payload?.subscriptionId
+
+      // Close connection
+      ws1.simulateClose(1006, 'Abnormal closure')
+      handler.handleClose(ws1 as any, 1006, 'Abnormal closure')
+
+      // Broadcast changes during reconnection window
+      handler.broadcastChange({
+        table: 'users',
+        operation: 'INSERT',
+        newRow: { id: 1, name: 'Alice' },
+      })
+
+      handler.broadcastChange({
+        table: 'users',
+        operation: 'INSERT',
+        newRow: { id: 2, name: 'Bob' },
+      })
+
+      // Reconnect
+      const ws2 = createMockWebSocket('ws-2')
+      safeHandleUpgrade(handler, mockState, ws2)
+      ws2.send.mockClear()
+
+      await handler.handleMessage(ws2 as any, JSON.stringify({
+        type: 'reconnect',
+        id: 'reconnect-1',
+        payload: {
+          connectionId: connectionId,
+          subscriptionIds: [subscriptionId],
+        },
+      }))
+
+      // Should receive buffered changes
+      // Note: This test may fail if buffering is not implemented correctly
+      const messages = getSentMessages(ws2)
+      expect(messages.length).toBeGreaterThan(0)
+    })
+
+    it('should handle concurrent reconnect attempts', async () => {
+      /**
+       * FLAKY SCENARIO: Multiple reconnect attempts arrive simultaneously,
+       * potentially causing state corruption.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      // Initial connection
+      const ws1 = createMockWebSocket('ws-1')
+      safeHandleUpgrade(handler, mockState, ws1)
+
+      const welcome1 = getLastSentMessage(ws1) as any
+      const connectionId = welcome1?.payload?.connectionId
+
+      // Close first connection
+      ws1.simulateClose(1006, 'Abnormal closure')
+      handler.handleClose(ws1 as any, 1006, 'Abnormal closure')
+
+      // Multiple concurrent reconnect attempts
+      const reconnectPromises: Promise<void>[] = []
+
+      for (let i = 0; i < 3; i++) {
+        const ws = createMockWebSocket(`ws-reconnect-${i}`)
+        safeHandleUpgrade(handler, mockState, ws)
+
+        reconnectPromises.push(handler.handleMessage(ws as any, JSON.stringify({
+          type: 'reconnect',
+          id: `reconnect-${i}`,
+          payload: { connectionId },
+        })))
+      }
+
+      // All should complete without errors
+      await expect(Promise.all(reconnectPromises)).resolves.not.toThrow()
+
+      // Handler should be in a consistent state
+      expect(handler.getConnectedCount()).toBeGreaterThan(0)
+    })
+  })
+
+  // ==========================================================================
+  // 3. Message Ordering Guarantees
+  // ==========================================================================
+
+  describe('Message Ordering Guarantees', () => {
+    it('should deliver changes in order', async () => {
+      /**
+       * FLAKY SCENARIO: Changes should be delivered in the order they occurred,
+       * but async processing might reorder them.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Subscribe
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      ws.send.mockClear()
+
+      // Broadcast multiple changes rapidly
+      const changeCount = 10
+      for (let i = 0; i < changeCount; i++) {
+        handler.broadcastChange({
+          table: 'users',
+          operation: 'INSERT',
+          newRow: { id: i, seq: i },
+        })
+      }
+
+      // Verify order
+      const messages = getSentMessages(ws)
+      const changeMessages = messages.filter(m => m.type === 'change')
+
+      expect(changeMessages.length).toBe(changeCount)
+
+      // Changes should be in order
+      for (let i = 0; i < changeMessages.length; i++) {
+        const payload = changeMessages[i].payload as { newRow: { seq: number } }
+        expect(payload.newRow.seq).toBe(i)
+      }
+    })
+
+    it('should not duplicate messages on slow consumers', async () => {
+      /**
+       * FLAKY SCENARIO: If a consumer is slow, messages might be sent multiple times.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Subscribe
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      ws.send.mockClear()
+
+      // Simulate slow send by making it take time
+      let sendCount = 0
+      ws.send.mockImplementation(() => {
+        sendCount++
+        // Simulate slow processing
+      })
+
+      // Broadcast change
+      handler.broadcastChange({
+        table: 'users',
+        operation: 'INSERT',
+        newRow: { id: 1 },
+      })
+
+      // Small delay
+      await flushPromises()
+
+      // Broadcast another change
+      handler.broadcastChange({
+        table: 'users',
+        operation: 'INSERT',
+        newRow: { id: 2 },
+      })
+
+      // Should only have 2 sends, no duplicates
+      expect(sendCount).toBe(2)
+    })
+
+    it('should maintain ordering across multiple subscribers', async () => {
+      /**
+       * FLAKY SCENARIO: Multiple subscribers should all receive messages in
+       * the same order.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      // Create multiple WebSocket connections
+      const webSockets: MockWebSocket[] = []
+      for (let i = 0; i < 3; i++) {
+        const ws = createMockWebSocket(`ws-${i}`)
+        safeHandleUpgrade(handler, mockState, ws)
+
+        // Subscribe
+        await handler.handleMessage(ws as any, JSON.stringify({
+          type: 'subscribe',
+          id: `sub-${i}`,
+          payload: { table: 'users' },
+        }))
+
+        ws.send.mockClear()
+        webSockets.push(ws)
+      }
+
+      // Broadcast changes
+      const changeCount = 5
+      for (let i = 0; i < changeCount; i++) {
+        handler.broadcastChange({
+          table: 'users',
+          operation: 'INSERT',
+          newRow: { id: i, seq: i },
+        })
+      }
+
+      // All subscribers should have received changes in the same order
+      const allOrders: number[][] = webSockets.map(ws => {
+        const messages = getSentMessages(ws)
+        const changeMessages = messages.filter(m => m.type === 'change')
+        return changeMessages.map(m => {
+          const payload = m.payload as { newRow: { seq: number } }
+          return payload.newRow.seq
+        })
+      })
+
+      // All orders should be identical
+      const expectedOrder = Array.from({ length: changeCount }, (_, i) => i)
+      for (const order of allOrders) {
+        expect(order).toEqual(expectedOrder)
+      }
+    })
+  })
+
+  // ==========================================================================
+  // 4. Connection State Management
+  // ==========================================================================
+
+  describe('Connection State Management', () => {
+    it('should handle WebSocket close during subscription', async () => {
+      /**
+       * FLAKY SCENARIO: WebSocket closes exactly while a subscription is being
+       * processed, leaving orphaned state.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Start subscription
+      const subscribePromise = handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      // Close during subscription processing
+      ws.simulateClose(1006, 'Abnormal closure')
+      handler.handleClose(ws as any, 1006, 'Abnormal closure')
+
+      // Should complete without errors
+      await subscribePromise
+
+      // State should be clean
+      expect(handler.getConnectedCount()).toBe(0)
+      expect(handler.getSubscriptionCount()).toBe(0)
+    })
+
+    it('should clean up all resources on connection close', async () => {
+      /**
+       * FLAKY SCENARIO: Resources (timers, subscriptions, state) should be
+       * fully cleaned up, even under concurrent operations.
+       */
+      timers = vi.useFakeTimers()
+
+      const config: CDCWebSocketConfig = {
+        heartbeatInterval: 100,
+        heartbeatTimeout: 50,
+      }
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, undefined, config)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Create subscriptions
+      for (let i = 0; i < 5; i++) {
+        await handler.handleMessage(ws as any, JSON.stringify({
+          type: 'subscribe',
+          id: `sub-${i}`,
+          payload: { table: `table-${i}` },
+        }))
+      }
+
+      expect(handler.getSubscriptionCount()).toBe(5)
+
+      // Close connection
+      ws.simulateClose(1000, 'Normal closure')
+      handler.handleClose(ws as any, 1000, 'Normal closure')
+
+      // All resources should be cleaned up
+      expect(handler.getSubscriptionCount()).toBe(0)
+      expect(handler.getConnectedCount()).toBe(0)
+
+      // Heartbeat timer should not fire after close
+      const sendCallsBefore = ws.send.mock.calls.length
+      await advanceTimersAndFlush(200)
+      const sendCallsAfter = ws.send.mock.calls.length
+
+      expect(sendCallsAfter).toBe(sendCallsBefore)
+    })
+
+    it('should handle error during message send', async () => {
+      /**
+       * FLAKY SCENARIO: If send() throws, should not corrupt connection state.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Subscribe
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      ws.send.mockClear()
+
+      // Make send throw
+      ws.send.mockImplementation(() => {
+        throw new Error('Network error')
+      })
+
+      // Should not throw
+      expect(() => {
+        handler.broadcastChange({
+          table: 'users',
+          operation: 'INSERT',
+          newRow: { id: 1 },
+        })
+      }).not.toThrow()
+
+      // Handler should still be functional
+      expect(handler.getConnectedCount()).toBe(1)
+    })
+
+    it('should handle multiple connections from same client', async () => {
+      /**
+       * FLAKY SCENARIO: Same client opens multiple connections, causing
+       * potential state confusion.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const connections: MockWebSocket[] = []
+      const connectionIds: string[] = []
+
+      // Open multiple connections
+      for (let i = 0; i < 3; i++) {
+        const ws = createMockWebSocket(`ws-same-client-${i}`)
+        safeHandleUpgrade(handler, mockState, ws)
+
+        const welcome = getLastSentMessage(ws) as any
+        connectionIds.push(welcome?.payload?.connectionId)
+        connections.push(ws)
+      }
+
+      // Each connection should have unique ID
+      expect(new Set(connectionIds).size).toBe(3)
+
+      // Subscribe each to different tables
+      for (let i = 0; i < connections.length; i++) {
+        await handler.handleMessage(connections[i] as any, JSON.stringify({
+          type: 'subscribe',
+          id: `sub-${i}`,
+          payload: { table: `table-${i}` },
+        }))
+      }
+
+      expect(handler.getSubscriptionCount()).toBe(3)
+
+      // Close one connection - should not affect others
+      connections[1].simulateClose(1000, 'Normal closure')
+      handler.handleClose(connections[1] as any, 1000, 'Normal closure')
+
+      expect(handler.getConnectedCount()).toBe(2)
+      expect(handler.getSubscriptionCount()).toBe(2)
+    })
+
+    it('should timeout queued operations if reconnection fails', async () => {
+      /**
+       * FLAKY SCENARIO: Operations queued during reconnection should timeout
+       * if reconnection ultimately fails.
+       *
+       * This is a known flaky test - stream closure is not handled properly.
+       */
+      timers = vi.useFakeTimers()
+
+      const config: CDCWebSocketConfig = {
+        reconnectDelay: 100,
+        maxReconnectAttempts: 2,
+      }
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, undefined, config)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+
+      const welcome = getLastSentMessage(ws) as any
+      const connectionId = welcome?.payload?.connectionId
+
+      // Subscribe
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      // Close connection
+      ws.simulateClose(1006, 'Abnormal closure')
+      handler.handleClose(ws as any, 1006, 'Abnormal closure')
+
+      // Wait for reconnection window to expire
+      await advanceTimersAndFlush(config.reconnectDelay! * (config.maxReconnectAttempts! + 1))
+
+      // Connection state should be fully cleaned up
+      expect(handler.getConnectionInfo(connectionId)).toBeNull()
+    })
+  })
+
+  // ==========================================================================
+  // 5. Edge Cases
+  // ==========================================================================
+
+  describe('Edge Cases', () => {
+    it('should handle max subscriptions limit', async () => {
+      /**
+       * Test that subscription limits are enforced.
+       */
+      const config: CDCWebSocketConfig = {
+        maxSubscriptions: 5,
+      }
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any, undefined, config)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+      ws.send.mockClear()
+
+      // Create max subscriptions
+      for (let i = 0; i < 5; i++) {
+        await handler.handleMessage(ws as any, JSON.stringify({
+          type: 'subscribe',
+          id: `sub-${i}`,
+          payload: { table: `table-${i}` },
+        }))
+      }
+
+      const messagesBeforeLimit = getSentMessages(ws)
+      const subscribedCount = messagesBeforeLimit.filter(m => m.type === 'subscribed').length
+      expect(subscribedCount).toBe(5)
+
+      ws.send.mockClear()
+
+      // Try to exceed limit
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-excess',
+        payload: { table: 'excess-table' },
+      }))
+
+      const messagesAfterLimit = getSentMessages(ws)
+      const errorMsg = messagesAfterLimit.find(m => m.type === 'error')
+      expect(errorMsg).toBeDefined()
+      expect((errorMsg?.payload as any)?.code).toBe('LIMIT_EXCEEDED')
+    })
+
+    it('should reject invalid message payloads gracefully', async () => {
+      /**
+       * Test that malformed messages are handled gracefully.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      safeHandleUpgrade(handler, mockState, ws)
+      ws.send.mockClear()
+
+      // Missing required fields
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'bad-sub',
+        payload: {}, // Missing 'table'
+      }))
+
+      const messages = getSentMessages(ws)
+      const errorMsg = messages.find(m => m.type === 'error')
+      expect(errorMsg).toBeDefined()
+    })
+
+    it('should handle WebSocket in CONNECTING state', async () => {
+      /**
+       * FLAKY SCENARIO: Operations arrive before WebSocket is fully open.
+       */
+      handler = createCDCWebSocketHandler(mockState as any, mockExecutor as any)
+
+      const ws = createMockWebSocket()
+      ws.setReadyState(WS_CONNECTING)
+      safeHandleUpgrade(handler, mockState, ws)
+
+      // Welcome message should not be sent to CONNECTING socket
+      // (depends on implementation - some might queue)
+
+      ws.setReadyState(WS_OPEN)
+
+      // Now should be able to subscribe
+      await handler.handleMessage(ws as any, JSON.stringify({
+        type: 'subscribe',
+        id: 'sub-1',
+        payload: { table: 'users' },
+      }))
+
+      // Should get subscribed response
+      const messages = getSentMessages(ws)
+      expect(messages.some(m => m.type === 'subscribed')).toBe(true)
+    })
+  })
+})