@comprehend/telemetry-node 0.1.4 → 0.2.1

package/src/ComprehendSDK.test.ts

@@ -0,0 +1,160 @@
+import { ComprehendSDK } from './ComprehendSDK';
+import { WebSocketConnection } from './WebSocketConnection';
+import { ComprehendDevSpanProcessor } from './ComprehendDevSpanProcessor';
+import { ComprehendMetricsExporter } from './ComprehendMetricsExporter';
+import { InitAck, CustomMetricSpecification } from './wire-protocol';
+
+jest.mock('./WebSocketConnection');
+jest.mock('./ComprehendDevSpanProcessor');
+jest.mock('./ComprehendMetricsExporter');
+
+const MockedWebSocketConnection = WebSocketConnection as jest.MockedClass<typeof WebSocketConnection>;
+const MockedSpanProcessor = ComprehendDevSpanProcessor as jest.MockedClass<typeof ComprehendDevSpanProcessor>;
+const MockedMetricsExporter = ComprehendMetricsExporter as jest.MockedClass<typeof ComprehendMetricsExporter>;
+
+describe('ComprehendSDK', () => {
+    let capturedOnAuthorized: ((ack: InitAck) => void) | undefined;
+    let capturedOnCustomMetricChange: ((specs: CustomMetricSpecification[]) => void) | undefined;
+    let mockConnectionInstance: any;
+
+    beforeEach(() => {
+        jest.clearAllMocks();
+        capturedOnAuthorized = undefined;
+        capturedOnCustomMetricChange = undefined;
+        mockConnectionInstance = null;
+
+        MockedWebSocketConnection.mockImplementation((options: any) => {
+            capturedOnAuthorized = options.onAuthorized;
+            capturedOnCustomMetricChange = options.onCustomMetricChange;
+            mockConnectionInstance = {
+                sendMessage: jest.fn(),
+                setProcessContext: jest.fn(),
+                nextSeq: jest.fn(),
+                close: jest.fn(),
+            };
+            return mockConnectionInstance;
+        });
+
+        MockedSpanProcessor.mockImplementation(() => ({
+            onStart: jest.fn(),
+            onEnd: jest.fn(),
+            forceFlush: jest.fn().mockResolvedValue(undefined),
+            shutdown: jest.fn().mockResolvedValue(undefined),
+            updateCustomMetrics: jest.fn(),
+        } as any));
+
+        MockedMetricsExporter.mockImplementation(() => ({
+            export: jest.fn(),
+            forceFlush: jest.fn().mockResolvedValue(undefined),
+            shutdown: jest.fn().mockResolvedValue(undefined),
+            selectAggregationTemporality: jest.fn(),
+            updateCustomMetrics: jest.fn(),
+        } as any));
+    });
+
+    it('should create a WebSocketConnection with the correct options', () => {
+        new ComprehendSDK({ organization: 'test-org', token: 'test-token' });
+
+        expect(MockedWebSocketConnection).toHaveBeenCalledTimes(1);
+        const opts = MockedWebSocketConnection.mock.calls[0][0];
+        expect(opts.organization).toBe('test-org');
+        expect(opts.token).toBe('test-token');
+    });
+
+    it('should pass console.log as logger when debug is true', () => {
+        new ComprehendSDK({ organization: 'org', token: 'tok', debug: true });
+
+        const opts = MockedWebSocketConnection.mock.calls[0][0];
+        expect(opts.logger).toBe(console.log);
+    });
+
+    it('should pass custom logger function', () => {
+        const myLogger = jest.fn();
+        new ComprehendSDK({ organization: 'org', token: 'tok', debug: myLogger });
+
+        const opts = MockedWebSocketConnection.mock.calls[0][0];
+        expect(opts.logger).toBe(myLogger);
+    });
+
+    it('should return the same span processor on multiple calls', () => {
+        const sdk = new ComprehendSDK({ organization: 'org', token: 'tok' });
+
+        const proc1 = sdk.getSpanProcessor();
+        const proc2 = sdk.getSpanProcessor();
+
+        expect(proc1).toBe(proc2);
+        expect(MockedSpanProcessor).toHaveBeenCalledTimes(1);
+    });
+
+    it('should return the same metrics exporter on multiple calls', () => {
+        const sdk = new ComprehendSDK({ organization: 'org', token: 'tok' });
+
+        const exp1 = sdk.getMetricsExporter();
+        const exp2 = sdk.getMetricsExporter();
+
+        expect(exp1).toBe(exp2);
+        expect(MockedMetricsExporter).toHaveBeenCalledTimes(1);
+    });
+
+    it('should share the same WebSocketConnection between processor and exporter', () => {
+        const sdk = new ComprehendSDK({ organization: 'org', token: 'tok' });
+
+        sdk.getSpanProcessor();
+        sdk.getMetricsExporter();
+
+        // Both should receive the same connection instance
+        expect(MockedSpanProcessor).toHaveBeenCalledWith(mockConnectionInstance);
+        expect(MockedMetricsExporter).toHaveBeenCalledWith(mockConnectionInstance);
+    });
+
+    it('should distribute custom metrics from onAuthorized to processor and exporter', () => {
+        const sdk = new ComprehendSDK({ organization: 'org', token: 'tok' });
+        const processor = sdk.getSpanProcessor();
+        const exporter = sdk.getMetricsExporter();
+
+        const specs: CustomMetricSpecification[] = [
+            { type: 'span', rule: { kind: 'type', value: 'server' }, subject: 's1' },
+            { type: 'cumulative', id: 'm1', attributes: ['a'], subject: 's2' },
+        ];
+
+        // Simulate the auth callback
+        capturedOnAuthorized!({ type: 'ack-authorized', customMetrics: specs });
+
+        expect((processor as any).updateCustomMetrics).toHaveBeenCalledWith(specs);
+        expect((exporter as any).updateCustomMetrics).toHaveBeenCalledWith(specs);
+    });
+
+    it('should distribute custom metrics from onCustomMetricChange', () => {
+        const sdk = new ComprehendSDK({ organization: 'org', token: 'tok' });
+        const processor = sdk.getSpanProcessor();
+        const exporter = sdk.getMetricsExporter();
+
+        const specs: CustomMetricSpecification[] = [
+            { type: 'timeseries', id: 'new-metric', attributes: [], subject: 's3' },
+        ];
+
+        capturedOnCustomMetricChange!(specs);
+
+        expect((processor as any).updateCustomMetrics).toHaveBeenCalledWith(specs);
+        expect((exporter as any).updateCustomMetrics).toHaveBeenCalledWith(specs);
+    });
+
+    it('should shutdown processor, exporter, and connection', async () => {
+        const sdk = new ComprehendSDK({ organization: 'org', token: 'tok' });
+        const processor = sdk.getSpanProcessor();
+        const exporter = sdk.getMetricsExporter();
+
+        await sdk.shutdown();
+
+        expect((processor as any).shutdown).toHaveBeenCalled();
+        expect((exporter as any).shutdown).toHaveBeenCalled();
+        expect(mockConnectionInstance.close).toHaveBeenCalled();
+    });
+
+    it('should shutdown cleanly even without processor or exporter created', async () => {
+        const sdk = new ComprehendSDK({ organization: 'org', token: 'tok' });
+
+        await expect(sdk.shutdown()).resolves.toBeUndefined();
+        expect(mockConnectionInstance.close).toHaveBeenCalled();
+    });
+});

package/src/ComprehendSDK.ts

@@ -0,0 +1,63 @@
+import { CustomMetricSpecification, InitAck } from './wire-protocol';
+import { WebSocketConnection } from './WebSocketConnection';
+import { ComprehendDevSpanProcessor } from './ComprehendDevSpanProcessor';
+import { ComprehendMetricsExporter } from './ComprehendMetricsExporter';
+
+export type LogFn = (message: string) => void;
+
+export class ComprehendSDK {
+    private readonly connection: WebSocketConnection;
+    private spanProcessor: ComprehendDevSpanProcessor | null = null;
+    private metricsExporter: ComprehendMetricsExporter | null = null;
+    private customMetricListeners: Array<(specs: CustomMetricSpecification[]) => void> = [];
+
+    constructor(options: { organization: string, token: string, debug?: boolean | LogFn }) {
+        const logger = options.debug === true ? console.log
+                      : options.debug === false ? undefined
+                      : options.debug;
+
+        this.connection = new WebSocketConnection({
+            organization: options.organization,
+            token: options.token,
+            logger,
+            onAuthorized: (ack: InitAck) => {
+                this.distributeCustomMetrics(ack.customMetrics);
+            },
+            onCustomMetricChange: (specs: CustomMetricSpecification[]) => {
+                this.distributeCustomMetrics(specs);
+            },
+        });
+    }
+
+    private distributeCustomMetrics(specs: CustomMetricSpecification[]): void {
+        for (const listener of this.customMetricListeners) {
+            listener(specs);
+        }
+    }
+
+    getSpanProcessor(): ComprehendDevSpanProcessor {
+        if (!this.spanProcessor) {
+            this.spanProcessor = new ComprehendDevSpanProcessor(this.connection);
+            this.customMetricListeners.push((specs) => this.spanProcessor!.updateCustomMetrics(specs));
+        }
+        return this.spanProcessor;
+    }
+
+    getMetricsExporter(): ComprehendMetricsExporter {
+        if (!this.metricsExporter) {
+            this.metricsExporter = new ComprehendMetricsExporter(this.connection);
+            this.customMetricListeners.push((specs) => this.metricsExporter!.updateCustomMetrics(specs));
+        }
+        return this.metricsExporter;
+    }
+
+    async shutdown(): Promise<void> {
+        if (this.spanProcessor) {
+            await this.spanProcessor.shutdown();
+        }
+        if (this.metricsExporter) {
+            await this.metricsExporter.shutdown();
+        }
+        this.connection.close();
+    }
+}