diff options
Diffstat (limited to 'packages/core/src/telemetry/metrics.test.ts')
| -rw-r--r-- | packages/core/src/telemetry/metrics.test.ts | 198 |
1 files changed, 163 insertions, 35 deletions
diff --git a/packages/core/src/telemetry/metrics.test.ts b/packages/core/src/telemetry/metrics.test.ts index 7e24b9ad..4fcdd9e1 100644 --- a/packages/core/src/telemetry/metrics.test.ts +++ b/packages/core/src/telemetry/metrics.test.ts @@ -5,32 +5,84 @@ */ import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; -import { Counter, Meter, metrics } from '@opentelemetry/api'; -import { initializeMetrics, recordTokenUsageMetrics } from './metrics.js'; +import type { + Counter, + Meter, + Attributes, + Context, + Histogram, +} from '@opentelemetry/api'; import { Config } from '../config/config.js'; +import { FileOperation } from './metrics.js'; -const mockCounter = { - add: vi.fn(), +const mockCounterAddFn: Mock< + (value: number, attributes?: Attributes, context?: Context) => void +> = vi.fn(); +const mockHistogramRecordFn: Mock< + (value: number, attributes?: Attributes, context?: Context) => void +> = vi.fn(); + +const mockCreateCounterFn: Mock<(name: string, options?: unknown) => Counter> = + vi.fn(); +const mockCreateHistogramFn: Mock< + (name: string, options?: unknown) => Histogram +> = vi.fn(); + +const mockCounterInstance = { + add: mockCounterAddFn, } as unknown as Counter; -const mockMeter = { - createCounter: vi.fn().mockReturnValue(mockCounter), - createHistogram: vi.fn().mockReturnValue({ record: vi.fn() }), +const mockHistogramInstance = { + record: mockHistogramRecordFn, +} as unknown as Histogram; + +const mockMeterInstance = { + createCounter: mockCreateCounterFn.mockReturnValue(mockCounterInstance), + createHistogram: mockCreateHistogramFn.mockReturnValue(mockHistogramInstance), } as unknown as Meter; -vi.mock('@opentelemetry/api', () => ({ - metrics: { - getMeter: vi.fn(), - }, - ValueType: { - INT: 1, - }, -})); +function originalOtelMockFactory() { + return { + metrics: { + getMeter: vi.fn(), + }, + ValueType: { + INT: 1, + }, + }; +} + +vi.mock('@opentelemetry/api', originalOtelMockFactory); describe('Telemetry Metrics', () => { - beforeEach(() => { - vi.clearAllMocks(); - (metrics.getMeter as Mock).mockReturnValue(mockMeter); + let initializeMetricsModule: typeof import('./metrics.js').initializeMetrics; + let recordTokenUsageMetricsModule: typeof import('./metrics.js').recordTokenUsageMetrics; + let recordFileOperationMetricModule: typeof import('./metrics.js').recordFileOperationMetric; + + beforeEach(async () => { + vi.resetModules(); + vi.doMock('@opentelemetry/api', () => { + const actualApi = originalOtelMockFactory(); + (actualApi.metrics.getMeter as Mock).mockReturnValue(mockMeterInstance); + return actualApi; + }); + + const metricsJsModule = await import('./metrics.js'); + initializeMetricsModule = metricsJsModule.initializeMetrics; + recordTokenUsageMetricsModule = metricsJsModule.recordTokenUsageMetrics; + recordFileOperationMetricModule = metricsJsModule.recordFileOperationMetric; + + const otelApiModule = await import('@opentelemetry/api'); + + mockCounterAddFn.mockClear(); + mockCreateCounterFn.mockClear(); + mockCreateHistogramFn.mockClear(); + mockHistogramRecordFn.mockClear(); + (otelApiModule.metrics.getMeter as Mock).mockClear(); + + (otelApiModule.metrics.getMeter as Mock).mockReturnValue(mockMeterInstance); + mockCreateCounterFn.mockReturnValue(mockCounterInstance); + mockCreateHistogramFn.mockReturnValue(mockHistogramInstance); }); describe('recordTokenUsageMetrics', () => { @@ -39,14 +91,18 @@ describe('Telemetry Metrics', () => { } as unknown as Config; it('should not record metrics if not initialized', () => { - recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input'); - expect(mockCounter.add).not.toHaveBeenCalled(); + recordTokenUsageMetricsModule(mockConfig, 'gemini-pro', 100, 'input'); + expect(mockCounterAddFn).not.toHaveBeenCalled(); }); it('should record token usage with the correct attributes', () => { - initializeMetrics(mockConfig); - recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input'); - expect(mockCounter.add).toHaveBeenCalledWith(100, { + initializeMetricsModule(mockConfig); + recordTokenUsageMetricsModule(mockConfig, 'gemini-pro', 100, 'input'); + expect(mockCounterAddFn).toHaveBeenCalledTimes(2); + expect(mockCounterAddFn).toHaveBeenNthCalledWith(1, 1, { + 'session.id': 'test-session-id', + }); + expect(mockCounterAddFn).toHaveBeenNthCalledWith(2, 100, { 'session.id': 'test-session-id', model: 'gemini-pro', type: 'input', @@ -54,30 +110,32 @@ describe('Telemetry Metrics', () => { }); it('should record token usage for different types', () => { - initializeMetrics(mockConfig); - recordTokenUsageMetrics(mockConfig, 'gemini-pro', 50, 'output'); - expect(mockCounter.add).toHaveBeenCalledWith(50, { + initializeMetricsModule(mockConfig); + mockCounterAddFn.mockClear(); + + recordTokenUsageMetricsModule(mockConfig, 'gemini-pro', 50, 'output'); + expect(mockCounterAddFn).toHaveBeenCalledWith(50, { 'session.id': 'test-session-id', model: 'gemini-pro', type: 'output', }); - recordTokenUsageMetrics(mockConfig, 'gemini-pro', 25, 'thought'); - expect(mockCounter.add).toHaveBeenCalledWith(25, { + recordTokenUsageMetricsModule(mockConfig, 'gemini-pro', 25, 'thought'); + expect(mockCounterAddFn).toHaveBeenCalledWith(25, { 'session.id': 'test-session-id', model: 'gemini-pro', type: 'thought', }); - recordTokenUsageMetrics(mockConfig, 'gemini-pro', 75, 'cache'); - expect(mockCounter.add).toHaveBeenCalledWith(75, { + recordTokenUsageMetricsModule(mockConfig, 'gemini-pro', 75, 'cache'); + expect(mockCounterAddFn).toHaveBeenCalledWith(75, { 'session.id': 'test-session-id', model: 'gemini-pro', type: 'cache', }); - recordTokenUsageMetrics(mockConfig, 'gemini-pro', 125, 'tool'); - expect(mockCounter.add).toHaveBeenCalledWith(125, { + recordTokenUsageMetricsModule(mockConfig, 'gemini-pro', 125, 'tool'); + expect(mockCounterAddFn).toHaveBeenCalledWith(125, { 'session.id': 'test-session-id', model: 'gemini-pro', type: 'tool', @@ -85,13 +143,83 @@ describe('Telemetry Metrics', () => { }); it('should handle different models', () => { - initializeMetrics(mockConfig); - recordTokenUsageMetrics(mockConfig, 'gemini-ultra', 200, 'input'); - expect(mockCounter.add).toHaveBeenCalledWith(200, { + initializeMetricsModule(mockConfig); + mockCounterAddFn.mockClear(); + + recordTokenUsageMetricsModule(mockConfig, 'gemini-ultra', 200, 'input'); + expect(mockCounterAddFn).toHaveBeenCalledWith(200, { 'session.id': 'test-session-id', model: 'gemini-ultra', type: 'input', }); }); }); + + describe('recordFileOperationMetric', () => { + const mockConfig = { + getSessionId: () => 'test-session-id', + } as unknown as Config; + + it('should not record metrics if not initialized', () => { + recordFileOperationMetricModule( + mockConfig, + FileOperation.CREATE, + 10, + 'text/plain', + 'txt', + ); + expect(mockCounterAddFn).not.toHaveBeenCalled(); + }); + + it('should record file creation with all attributes', () => { + initializeMetricsModule(mockConfig); + recordFileOperationMetricModule( + mockConfig, + FileOperation.CREATE, + 10, + 'text/plain', + 'txt', + ); + + expect(mockCounterAddFn).toHaveBeenCalledTimes(2); + expect(mockCounterAddFn).toHaveBeenNthCalledWith(1, 1, { + 'session.id': 'test-session-id', + }); + expect(mockCounterAddFn).toHaveBeenNthCalledWith(2, 1, { + 'session.id': 'test-session-id', + operation: FileOperation.CREATE, + lines: 10, + mimetype: 'text/plain', + extension: 'txt', + }); + }); + + it('should record file read with minimal attributes', () => { + initializeMetricsModule(mockConfig); + mockCounterAddFn.mockClear(); + + recordFileOperationMetricModule(mockConfig, FileOperation.READ); + expect(mockCounterAddFn).toHaveBeenCalledWith(1, { + 'session.id': 'test-session-id', + operation: FileOperation.READ, + }); + }); + + it('should record file update with some attributes', () => { + initializeMetricsModule(mockConfig); + mockCounterAddFn.mockClear(); + + recordFileOperationMetricModule( + mockConfig, + FileOperation.UPDATE, + undefined, + 'application/javascript', + ); + expect(mockCounterAddFn).toHaveBeenCalledWith(1, { + 'session.id': 'test-session-id', + operation: FileOperation.UPDATE, + mimetype: 'application/javascript', + }); + }); + }); }); |