@comprehend/telemetry-node 0.1.4 → 0.2.1

package/src/WebSocketConnection.test.ts

@@ -5,12 +5,23 @@ import {
     NewObservedHttpRouteMessage,
     ObservationMessage,
     HttpServerObservation,
-    ObservationOutputMessage
+    ObservationOutputMessage,
+    TimeSeriesMetricsMessage,
+    CumulativeMetricsMessage,
+    TraceSpansMessage,
+    DatabaseQueryMessage,
+    InitAck,
+    CustomMetricSpecification,
 } from './wire-protocol';
 
 // Mock the ws WebSocket library
 jest.mock('ws');
 
+// Mock crypto.randomUUID
+jest.mock('crypto', () => ({
+    randomUUID: jest.fn().mockReturnValue('test-uuid-1234'),
+}));
+
 import WebSocket from 'ws';
 
 const MockedWebSocket = WebSocket as jest.MockedClass<typeof WebSocket>;
@@ -23,13 +34,16 @@ interface MockWebSocketInstance {
     readyState: number;
     OPEN: number;
     CLOSED: number;
-    // Event handlers that tests can trigger
     _triggerOpen: () => void;
     _triggerMessage: (data: string) => void;
     _triggerClose: (code: number, reason: string) => void;
     _triggerError: (error: Error) => void;
 }
 
+function authAck(customMetrics: CustomMetricSpecification[] = []): string {
+    return JSON.stringify({ type: 'ack-authorized', customMetrics } as InitAck);
+}
+
 describe('WebSocketConnection', () => {
     let mockSocket: MockWebSocketInstance;
     let connection: WebSocketConnection;
@@ -37,7 +51,6 @@ describe('WebSocketConnection', () => {
     let mockLogger: jest.Mock;
 
     beforeEach(() => {
-        // Reset all mocks
         jest.clearAllMocks();
         jest.clearAllTimers();
         jest.useFakeTimers();
@@ -47,7 +60,6 @@ describe('WebSocketConnection', () => {
             logMessages.push(message);
         });
 
-        // Create mock WebSocket instance with event handling
         mockSocket = {
             on: jest.fn(),
             send: jest.fn(),
@@ -55,13 +67,12 @@ describe('WebSocketConnection', () => {
             readyState: WebSocket.CLOSED,
             OPEN: WebSocket.OPEN,
             CLOSED: WebSocket.CLOSED,
-            _triggerOpen: function() { /* will be set up below */ },
-            _triggerMessage: function(data: string) { /* will be set up below */ },
-            _triggerClose: function(code: number, reason: string) { /* will be set up below */ },
-            _triggerError: function(error: Error) { /* will be set up below */ }
+            _triggerOpen: function() {},
+            _triggerMessage: function(data: string) {},
+            _triggerClose: function(code: number, reason: string) {},
+            _triggerError: function(error: Error) {}
         };
 
-        // Set up event handler storage and triggering
         const eventHandlers: Record<string, Function> = {};
         mockSocket.on.mockImplementation((event: string, handler: Function) => {
             eventHandlers[event] = handler;
@@ -83,7 +94,6 @@ describe('WebSocketConnection', () => {
             eventHandlers['error']?.(error);
         };
 
-        // Mock WebSocket constructor to return our mock instance
         MockedWebSocket.mockImplementation(() => mockSocket as any);
     });
 
@@ -96,7 +106,7 @@ describe('WebSocketConnection', () => {
 
     describe('Connection Establishment', () => {
         it('should create WebSocket connection with correct URL', () => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
 
             expect(MockedWebSocket).toHaveBeenCalledWith('wss://ingestion.comprehend.dev/test-org/observations');
             expect(mockSocket.on).toHaveBeenCalledWith('open', expect.any(Function));
@@ -105,21 +115,21 @@ describe('WebSocketConnection', () => {
             expect(mockSocket.on).toHaveBeenCalledWith('error', expect.any(Function));
         });
 
-        it('should send init message on connection open', () => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+        it('should send V2 init message on connection open', () => {
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
 
             mockSocket._triggerOpen();
 
             expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify({
                 event: 'init',
-                protocolVersion: 1,
+                protocolVersion: 2,
                 token: 'test-token'
             }));
             expect(logMessages).toContain('WebSocket connected. Sending init/auth message.');
         });
 
         it('should work without logger', () => {
-            connection = new WebSocketConnection('test-org', 'test-token');
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token' });
 
             mockSocket._triggerOpen();
 
@@ -130,23 +140,31 @@ describe('WebSocketConnection', () => {
 
     describe('Authorization Flow', () => {
         beforeEach(() => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
             mockSocket._triggerOpen();
-            jest.clearAllMocks(); // Clear the init message send
+            jest.clearAllMocks();
         });
 
-        it('should handle authorization acknowledgment', () => {
-            const authAck: ObservationOutputMessage = {
-                type: 'ack-authorized'
-            };
-
-            mockSocket._triggerMessage(JSON.stringify(authAck));
+        it('should handle authorization acknowledgment with customMetrics', () => {
+            mockSocket._triggerMessage(authAck([{ type: 'cumulative', id: 'test', attributes: [], subject: 'sub' }]));
 
             expect(logMessages).toContain('Authorization acknowledged by server.');
         });
 
+        it('should call onAuthorized callback with InitAck', () => {
+            const onAuthorized = jest.fn();
+            connection = new WebSocketConnection({
+                organization: 'test-org', token: 'test-token', logger: mockLogger, onAuthorized
+            });
+            mockSocket._triggerOpen();
+
+            const customMetrics: CustomMetricSpecification[] = [{ type: 'cumulative', id: 'test', attributes: [], subject: 'sub' }];
+            mockSocket._triggerMessage(authAck(customMetrics));
+
+            expect(onAuthorized).toHaveBeenCalledWith({ type: 'ack-authorized', customMetrics });
+        });
+
         it('should replay queued messages after authorization', () => {
-            // Queue some messages before authorization
             const serviceMessage: NewObservedServiceMessage = {
                 event: 'new-entity',
                 type: 'service',
@@ -160,6 +178,8 @@ describe('WebSocketConnection', () => {
                 observations: [{
                     type: 'http-server',
                     subject: 'test-subject',
+                    spanId: 'span1',
+                    traceId: 'trace1',
                     timestamp: [1700000000, 0],
                     path: '/test',
                     status: 200,
@@ -170,21 +190,17 @@ describe('WebSocketConnection', () => {
             connection.sendMessage(serviceMessage);
             connection.sendMessage(observationMessage);
 
-            // Should not send immediately (not authorized yet)
             expect(mockSocket.send).not.toHaveBeenCalled();
 
-            // Now authorize
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            mockSocket._triggerMessage(authAck());
 
-            // Should replay both messages
             expect(mockSocket.send).toHaveBeenCalledTimes(2);
             expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(serviceMessage));
             expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(observationMessage));
         });
 
         it('should send messages immediately when already authorized', () => {
-            // First authorize
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            mockSocket._triggerMessage(authAck());
             jest.clearAllMocks();
 
             const serviceMessage: NewObservedServiceMessage = {
@@ -200,104 +216,251 @@ describe('WebSocketConnection', () => {
         });
     });
 
+    describe('Context Handling', () => {
+        beforeEach(() => {
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
+            mockSocket._triggerOpen();
+        });
+
+        it('should send context-start after auth when process context is set before auth', () => {
+            connection.setProcessContext('service-hash', { 'service.name': 'test' });
+
+            // Not authorized yet, so no context-start sent
+            mockSocket.send.mockClear();
+
+            mockSocket._triggerMessage(authAck());
+
+            // Should have sent context-start first
+            const sentMessages = mockSocket.send.mock.calls.map(c => JSON.parse(c[0]));
+            const contextMsg = sentMessages.find((m: any) => m.event === 'context-start');
+            expect(contextMsg).toMatchObject({
+                event: 'context-start',
+                type: 'process',
+                serviceEntityHash: 'service-hash',
+                resources: { 'service.name': 'test' },
+                ingestionId: expect.any(String),
+            });
+        });
+
+        it('should send context-start immediately when already authorized', () => {
+            mockSocket._triggerMessage(authAck());
+            mockSocket.send.mockClear();
+
+            connection.setProcessContext('service-hash', { 'service.name': 'test' });
+
+            const sentMessages = mockSocket.send.mock.calls.map(c => JSON.parse(c[0]));
+            expect(sentMessages[0]).toMatchObject({
+                event: 'context-start',
+                type: 'process',
+                serviceEntityHash: 'service-hash',
+            });
+        });
+
+        it('should re-send context-start on reconnect before other replays', () => {
+            connection.setProcessContext('service-hash', { 'service.name': 'test' });
+            mockSocket._triggerMessage(authAck());
+
+            // Send an entity message
+            const entityMsg: NewObservedServiceMessage = {
+                event: 'new-entity', type: 'service', hash: 'h1', name: 'svc'
+            };
+            connection.sendMessage(entityMsg);
+
+            // Reconnect
+            mockSocket._triggerClose(1000, 'test');
+            jest.advanceTimersByTime(1000);
+            mockSocket._triggerOpen();
+            mockSocket.send.mockClear();
+
+            mockSocket._triggerMessage(authAck());
+
+            const sentMessages = mockSocket.send.mock.calls.map(c => JSON.parse(c[0]));
+            // Context should come first
+            expect(sentMessages[0].event).toBe('context-start');
+        });
+    });
+
     describe('Message Acknowledgment Handling', () => {
         beforeEach(() => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
             mockSocket._triggerOpen();
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            mockSocket._triggerMessage(authAck());
             jest.clearAllMocks();
         });
 
         it('should handle entity/interaction acknowledgments', () => {
             const serviceMessage: NewObservedServiceMessage = {
-                event: 'new-entity',
-                type: 'service',
-                hash: 'service-hash',
-                name: 'test-service'
+                event: 'new-entity', type: 'service', hash: 'service-hash', name: 'test-service'
             };
-
             const routeMessage: NewObservedHttpRouteMessage = {
-                event: 'new-entity',
-                type: 'http-route',
-                hash: 'route-hash',
-                parent: 'service-hash',
-                method: 'GET',
-                route: '/test'
+                event: 'new-entity', type: 'http-route', hash: 'route-hash',
+                parent: 'service-hash', method: 'GET', route: '/test'
             };
 
-            // Send messages
             connection.sendMessage(serviceMessage);
             connection.sendMessage(routeMessage);
 
-            // Acknowledge the service message
-            mockSocket._triggerMessage(JSON.stringify({
-                type: 'ack-observed',
-                hash: 'service-hash'
-            }));
+            // Acknowledge the service
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-observed', hash: 'service-hash' }));
 
-            // Clear previous calls before reconnection test
             mockSocket.send.mockClear();
 
-            // If we reconnect, only the route message should be replayed
+            // Reconnect — only unacked route should replay
             mockSocket._triggerClose(1000, 'test close');
             jest.advanceTimersByTime(1000);
             mockSocket._triggerOpen();
-
-            // Clear the init message that gets sent on reconnection
             mockSocket.send.mockClear();
 
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            mockSocket._triggerMessage(authAck());
 
-            // Should only send the unacknowledged route message
             expect(mockSocket.send).toHaveBeenCalledTimes(1);
             expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(routeMessage));
         });
 
         it('should handle observation acknowledgments', () => {
-            const observationMessage1: ObservationMessage = {
-                event: 'observations',
-                seq: 1,
-                observations: []
-            };
+            const obs1: ObservationMessage = { event: 'observations', seq: 1, observations: [] };
+            const obs2: ObservationMessage = { event: 'observations', seq: 2, observations: [] };
 
-            const observationMessage2: ObservationMessage = {
-                event: 'observations',
-                seq: 2,
-                observations: []
+            connection.sendMessage(obs1);
+            connection.sendMessage(obs2);
+
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-observations', seq: 1 }));
+
+            mockSocket.send.mockClear();
+            mockSocket._triggerClose(1000, 'test');
+            jest.advanceTimersByTime(1000);
+            mockSocket._triggerOpen();
+            mockSocket.send.mockClear();
+
+            mockSocket._triggerMessage(authAck());
+
+            expect(mockSocket.send).toHaveBeenCalledTimes(1);
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(obs2));
+        });
+
+        it('should handle timeseries acknowledgments', () => {
+            const msg: TimeSeriesMetricsMessage = {
+                event: 'timeseries', seq: 5, data: []
             };
+            connection.sendMessage(msg);
 
-            // Send messages
-            connection.sendMessage(observationMessage1);
-            connection.sendMessage(observationMessage2);
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-timeseries', seq: 5 }));
 
-            // Acknowledge first observation
-            mockSocket._triggerMessage(JSON.stringify({
-                type: 'ack-observations',
-                seq: 1
-            }));
+            mockSocket.send.mockClear();
+            mockSocket._triggerClose(1000, 'test');
+            jest.advanceTimersByTime(1000);
+            mockSocket._triggerOpen();
+            mockSocket.send.mockClear();
+            mockSocket._triggerMessage(authAck());
+
+            // Should not replay the acked message
+            const sentEvents = mockSocket.send.mock.calls.map(c => JSON.parse(c[0]).event);
+            expect(sentEvents).not.toContain('timeseries');
+        });
+
+        it('should handle cumulative acknowledgments', () => {
+            const msg: CumulativeMetricsMessage = { event: 'cumulative', seq: 3, data: [] };
+            connection.sendMessage(msg);
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-cumulative', seq: 3 }));
 
-            // Clear previous calls before reconnection test
             mockSocket.send.mockClear();
+            mockSocket._triggerClose(1000, 'test');
+            jest.advanceTimersByTime(1000);
+            mockSocket._triggerOpen();
+            mockSocket.send.mockClear();
+            mockSocket._triggerMessage(authAck());
 
-            // Reconnect - only second observation should be replayed
-            mockSocket._triggerClose(1000, 'test close');
+            const sentEvents = mockSocket.send.mock.calls.map(c => JSON.parse(c[0]).event);
+            expect(sentEvents).not.toContain('cumulative');
+        });
+
+        it('should handle tracespans acknowledgments', () => {
+            const msg: TraceSpansMessage = { event: 'tracespans', seq: 7, data: [] };
+            connection.sendMessage(msg);
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-tracespans', seq: 7 }));
+
+            mockSocket.send.mockClear();
+            mockSocket._triggerClose(1000, 'test');
             jest.advanceTimersByTime(1000);
             mockSocket._triggerOpen();
+            mockSocket.send.mockClear();
+            mockSocket._triggerMessage(authAck());
+
+            const sentEvents = mockSocket.send.mock.calls.map(c => JSON.parse(c[0]).event);
+            expect(sentEvents).not.toContain('tracespans');
+        });
 
-            // Clear the init message that gets sent on reconnection
+        it('should handle db-query acknowledgments', () => {
+            const msg: DatabaseQueryMessage = {
+                event: 'db-query', seq: 9, query: 'SELECT 1', from: 'a', to: 'b',
+                timestamp: [0, 0], duration: [0, 0]
+            };
+            connection.sendMessage(msg);
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-db-query', seq: 9 }));
+
+            mockSocket.send.mockClear();
+            mockSocket._triggerClose(1000, 'test');
+            jest.advanceTimersByTime(1000);
+            mockSocket._triggerOpen();
             mockSocket.send.mockClear();
+            mockSocket._triggerMessage(authAck());
 
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            const sentEvents = mockSocket.send.mock.calls.map(c => JSON.parse(c[0]).event);
+            expect(sentEvents).not.toContain('db-query');
+        });
 
-            // Should only send the unacknowledged second observation message
-            expect(mockSocket.send).toHaveBeenCalledTimes(1);
-            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(observationMessage2));
+        it('should handle context acknowledgments', () => {
+            connection.setProcessContext('svc-hash', { 'service.name': 'test' });
+            mockSocket._triggerMessage(authAck());
+
+            // Find the context seq
+            const contextCall = mockSocket.send.mock.calls.find(c => {
+                const m = JSON.parse(c[0]);
+                return m.event === 'context-start';
+            });
+            const contextSeq = JSON.parse(contextCall![0]).seq;
+
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-context', seq: contextSeq }));
+
+            // On reconnect, context-start should still be sent (because setProcessContext data persists)
+            // but the acked one should be removed from the unack queue
+        });
+    });
+
+    describe('Custom Metric Change', () => {
+        it('should call onCustomMetricChange callback', () => {
+            const onCustomMetricChange = jest.fn();
+            connection = new WebSocketConnection({
+                organization: 'test-org', token: 'test-token', onCustomMetricChange
+            });
+            mockSocket._triggerOpen();
+            mockSocket._triggerMessage(authAck());
+
+            const specs: CustomMetricSpecification[] = [
+                { type: 'timeseries', id: 'metric1', attributes: ['a'], subject: 'sub1' }
+            ];
+            mockSocket._triggerMessage(JSON.stringify({
+                type: 'custom-metric-change',
+                customMetrics: specs
+            }));
+
+            expect(onCustomMetricChange).toHaveBeenCalledWith(specs);
+        });
+    });
+
+    describe('Shared Sequence Counter', () => {
+        it('should provide incrementing sequence numbers', () => {
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token' });
+
+            expect(connection.nextSeq()).toBe(1);
+            expect(connection.nextSeq()).toBe(2);
+            expect(connection.nextSeq()).toBe(3);
         });
     });
 
     describe('Reconnection Logic', () => {
         beforeEach(() => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
             mockSocket._triggerOpen();
         });
 
@@ -307,39 +470,26 @@ describe('WebSocketConnection', () => {
             mockSocket._triggerClose(1000, 'Normal closure');
 
             expect(logMessages).toContain('WebSocket disconnected. Code: 1000, Reason: Normal closure');
-
-            // Should schedule reconnection
             expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 1000);
 
-            // Advance timer to trigger reconnection
             jest.advanceTimersByTime(1000);
-
-            // Should create new WebSocket connection
             expect(MockedWebSocket).toHaveBeenCalledTimes(2);
 
             setTimeoutSpy.mockRestore();
         });
 
         it('should reset authorization state on close', () => {
-            // First authorize
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            mockSocket._triggerMessage(authAck());
 
-            // Send a message (should send immediately)
             const testMessage: NewObservedServiceMessage = {
-                event: 'new-entity',
-                type: 'service',
-                hash: 'test-hash',
-                name: 'test-service'
+                event: 'new-entity', type: 'service', hash: 'test-hash', name: 'test-service'
             };
             connection.sendMessage(testMessage);
             expect(mockSocket.send).toHaveBeenCalled();
 
             jest.clearAllMocks();
-
-            // Close connection
             mockSocket._triggerClose(1000, 'test');
 
-            // Send another message - should queue (not send immediately)
             connection.sendMessage(testMessage);
             expect(mockSocket.send).not.toHaveBeenCalled();
         });
@@ -348,15 +498,14 @@ describe('WebSocketConnection', () => {
             connection.close();
             mockSocket._triggerClose(1000, 'Explicit close');
 
-            // Should not schedule reconnection
             jest.advanceTimersByTime(5000);
-            expect(MockedWebSocket).toHaveBeenCalledTimes(1); // Only initial connection
+            expect(MockedWebSocket).toHaveBeenCalledTimes(1);
         });
     });
 
     describe('Error Handling', () => {
         beforeEach(() => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
         });
 
         it('should log WebSocket errors', () => {
@@ -370,7 +519,6 @@ describe('WebSocketConnection', () => {
             mockSocket._triggerOpen();
             mockSocket._triggerMessage('{invalid json}');
 
-            // Check that an error message was logged (exact message varies by Node version)
             const errorLogExists = logMessages.some(msg =>
                 msg.startsWith('Error parsing message from server:') &&
                 msg.includes('JSON')
@@ -378,158 +526,91 @@ describe('WebSocketConnection', () => {
             expect(errorLogExists).toBe(true);
         });
 
-        it('should handle non-string error objects in message parsing', () => {
-            mockSocket._triggerOpen();
-
-            // Mock JSON.parse to throw a non-Error object
-            const originalParse = JSON.parse;
-            JSON.parse = jest.fn().mockImplementation(() => {
-                throw 'string error';
-            });
-
-            mockSocket._triggerMessage('{}');
-
-            expect(logMessages).toContain('Error parsing message from server: string error');
-
-            // Restore original JSON.parse
-            JSON.parse = originalParse;
-        });
-
         it('should not send messages when socket is not open', () => {
             mockSocket.readyState = WebSocket.CLOSED;
             mockSocket._triggerOpen();
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
-
-            // Clear the init and any previous messages
+            mockSocket._triggerMessage(authAck());
             mockSocket.send.mockClear();
 
-            // Force socket to closed state
             mockSocket.readyState = WebSocket.CLOSED;
 
             const testMessage: NewObservedServiceMessage = {
-                event: 'new-entity',
-                type: 'service',
-                hash: 'test-hash',
-                name: 'test-service'
+                event: 'new-entity', type: 'service', hash: 'test-hash', name: 'test-service'
             };
-
             connection.sendMessage(testMessage);
 
-            // Should not call send when socket is closed
             expect(mockSocket.send).not.toHaveBeenCalled();
         });
     });
 
     describe('Connection Lifecycle', () => {
         it('should close WebSocket connection properly', () => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
-
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
             connection.close();
-
             expect(mockSocket.close).toHaveBeenCalled();
         });
 
         it('should handle close when socket is null', () => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
-
-            // Simulate socket being null
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
             (connection as any).socket = null;
-
             expect(() => connection.close()).not.toThrow();
         });
     });
 
-    describe('Message Types and Queuing', () => {
+    describe('Message Queuing for All Types', () => {
         beforeEach(() => {
-            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            connection = new WebSocketConnection({ organization: 'test-org', token: 'test-token', logger: mockLogger });
             mockSocket._triggerOpen();
-        });
-
-        it('should queue new-entity messages', () => {
-            // Clear the initial init message
             mockSocket.send.mockClear();
+        });
 
-            const serviceMessage: NewObservedServiceMessage = {
-                event: 'new-entity',
-                type: 'service',
-                hash: 'service-hash',
-                name: 'test-service'
+        it('should queue and replay timeseries messages', () => {
+            const msg: TimeSeriesMetricsMessage = {
+                event: 'timeseries', seq: 1, data: [
+                    { subject: 's', type: 't', timestamp: [0, 0], value: 42, unit: 'By', attributes: {} }
+                ]
             };
-
-            connection.sendMessage(serviceMessage);
-
-            // Should be queued, not sent immediately (not authorized)
+            connection.sendMessage(msg);
             expect(mockSocket.send).not.toHaveBeenCalled();
 
-            // Authorize and check replay
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
-            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(serviceMessage));
+            mockSocket._triggerMessage(authAck());
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(msg));
         });
 
-        it('should queue new-interaction messages', () => {
-            // Clear the initial init message
-            mockSocket.send.mockClear();
-
-            const httpRequestMessage = {
-                event: 'new-interaction' as const,
-                type: 'http-request' as const,
-                hash: 'request-hash',
-                from: 'service-hash',
-                to: 'http-service-hash'
+        it('should queue and replay cumulative messages', () => {
+            const msg: CumulativeMetricsMessage = {
+                event: 'cumulative', seq: 2, data: []
             };
-
-            connection.sendMessage(httpRequestMessage);
-
-            // Should be queued
+            connection.sendMessage(msg);
             expect(mockSocket.send).not.toHaveBeenCalled();
 
-            // Authorize and check replay
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
-            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(httpRequestMessage));
+            mockSocket._triggerMessage(authAck());
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(msg));
         });
 
-        it('should queue observation messages', () => {
-            // Clear the initial init message
-            mockSocket.send.mockClear();
-
-            const observationMessage: ObservationMessage = {
-                event: 'observations',
-                seq: 1,
-                observations: [{
-                    type: 'http-server',
-                    subject: 'route-hash',
-                    timestamp: [1700000000, 0],
-                    path: '/test',
-                    status: 200,
-                    duration: [0, 50000000]
-                } as HttpServerObservation]
+        it('should queue and replay tracespans messages', () => {
+            const msg: TraceSpansMessage = {
+                event: 'tracespans', seq: 3, data: [
+                    { trace: 't1', span: 's1', parent: '', name: 'test', timestamp: [0, 0] }
+                ]
             };
-
-            connection.sendMessage(observationMessage);
-
-            // Should be queued
+            connection.sendMessage(msg);
             expect(mockSocket.send).not.toHaveBeenCalled();
 
-            // Authorize and check replay
-            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
-            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(observationMessage));
+            mockSocket._triggerMessage(authAck());
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(msg));
         });
 
-        it('should not queue init messages in unacknowledged maps', () => {
-            // This is implicit - init messages are sent via sendRaw, not sendMessage
-            // and don't go through the queuing logic
-
-            const initMessage: InitMessage = {
-                event: 'init',
-                protocolVersion: 1,
-                token: 'test-token'
+        it('should queue and replay db-query messages', () => {
+            const msg: DatabaseQueryMessage = {
+                event: 'db-query', seq: 4, query: 'SELECT 1', from: 'a', to: 'b',
+                timestamp: [0, 0], duration: [0, 1000000]
             };
+            connection.sendMessage(msg);
+            expect(mockSocket.send).not.toHaveBeenCalled();
 
-            // Manually test sendRaw behavior (private method testing)
-            mockSocket.readyState = WebSocket.OPEN;
-            (connection as any).sendRaw(initMessage);
-
-            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(initMessage));
+            mockSocket._triggerMessage(authAck());
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(msg));
         });
     });
-});
+});