From 72d55a5b4077b3387a340c29b78fe48b37140d4d Mon Sep 17 00:00:00 2001 From: Stefano Date: Thu, 17 Jul 2025 11:35:22 -0500 Subject: [PATCH] Added unit tests for posthog-llm.tool.ts --- src/tools/posthog-llm.tool.test.ts | 433 +++++++++++++++++++++++++++++ src/tools/posthog-llm.tool.ts | 10 +- src/utils/__mocks__/logger.util.ts | 18 ++ src/utils/logger.util.ts | 4 +- 4 files changed, 457 insertions(+), 8 deletions(-) create mode 100644 src/tools/posthog-llm.tool.test.ts create mode 100644 src/utils/__mocks__/logger.util.ts diff --git a/src/tools/posthog-llm.tool.test.ts b/src/tools/posthog-llm.tool.test.ts new file mode 100644 index 0000000..c5675d8 --- /dev/null +++ b/src/tools/posthog-llm.tool.test.ts @@ -0,0 +1,433 @@ +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import { Logger } from '../utils/logger.util.js'; +import { formatErrorForMcpTool } from '../utils/error.util.js'; +import { PostHogController } from '../controllers/posthog-llm.controller.js'; +import { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; +import { + PostHogLlmPropertiesPayloadSchema, + PostHogLlmPropertiesPayloadSchemaType, +} from '../types/posthog-llm.types.js'; +import posthogLlmTool from '../tools/posthog-llm.tool.js'; +import { mockDebug, mockError } from '../utils/__mocks__/logger.util'; + +// Mock all dependencies +jest.mock('../utils/logger.util.js', () => + jest.requireActual('../utils/__mocks__/logger.util.js'), +); +jest.mock('../utils/error.util.js'); +jest.mock('../controllers/posthog-llm.controller.js'); +jest.mock('../types/posthog-llm.types.js'); + +// Mock the capturePosthogLlmObservability function since it's not exported +// We'll need to access it through the server.tool registration +const mockCapturePosthogLlmObservability = jest.fn(); + +// Mock MCP Server +const mockMcpServer = { + tool: jest.fn(), +} as unknown as McpServer; + +describe('PostHog LLM Tool', () => { + let mockPostHogController: jest.Mocked; + let mockFormatErrorForMcpTool: jest.MockedFunction< + typeof formatErrorForMcpTool + >; + let mockPostHogLlmPropertiesPayloadSchema: jest.Mocked; + + beforeEach(() => { + // Reset all mocks + jest.clearAllMocks(); + mockDebug.mockClear(); + mockError.mockClear(); + + // Setup controller mock + mockPostHogController = { + capture: jest.fn(), + } as any; + (PostHogController as any).capture = mockPostHogController.capture; + + // Setup error formatter mock + mockFormatErrorForMcpTool = + formatErrorForMcpTool as jest.MockedFunction< + typeof formatErrorForMcpTool + >; + + // Setup schema mock + mockPostHogLlmPropertiesPayloadSchema = { + parse: jest.fn(), + shape: { + userId: { type: 'string' }, + model: { type: 'string' }, + provider: { type: 'string' }, + }, + }; + (PostHogLlmPropertiesPayloadSchema as any) = + mockPostHogLlmPropertiesPayloadSchema; + }); + + describe('registerTools', () => { + it('should register the llm_observability_posthog tool with correct parameters', () => { + // Act + posthogLlmTool.registerTools(mockMcpServer); + + // Assert + expect(mockMcpServer.tool).toHaveBeenCalledWith( + 'llm_observability_posthog', + 'Captures LLM usage in PostHog for observability, including requests, responses, and performance metrics', + mockPostHogLlmPropertiesPayloadSchema.shape, + expect.any(Function), + ); + expect(mockDebug).toHaveBeenCalledWith( + 'Registering PostHog LLM observability tools...', + ); + expect(mockDebug).toHaveBeenCalledWith( + 'Successfully registered llm_observability_posthog tool.', + ); + }); + + it('should create logger with correct context', () => { + // Act + posthogLlmTool.registerTools(mockMcpServer); + + // Assert + expect(Logger.forContext).toHaveBeenCalledWith( + 'tools/posthog-llm.tool.ts', + 'registerTools', + ); + }); + }); + + describe('capturePosthogLlmObservability', () => { + let captureFunction: ( + args: PostHogLlmPropertiesPayloadSchemaType, + ) => Promise; + + beforeEach(() => { + // Register the tool to get access to the capture function + posthogLlmTool.registerTools(mockMcpServer); + captureFunction = (mockMcpServer.tool as jest.Mock).mock + .calls[0][3]; + }); + + it('should successfully capture LLM observability data', async () => { + // Arrange + const inputArgs: PostHogLlmPropertiesPayloadSchemaType = { + userId: 'test-user-123', + model: 'gpt-4', + provider: 'openai', + input: 'test input', + outputChoices: ['choice1', 'choice2'], + traceId: 'trace-123', + inputTokens: 100, + outputTokens: 200, + latency: 1500, + httpStatus: 200, + baseUrl: 'https://api.openai.com', + }; + + const expectedParsedArgs = { ...inputArgs }; + const expectedControllerResponse = { + content: 'Successfully captured LLM observability data', + }; + + mockPostHogLlmPropertiesPayloadSchema.parse.mockReturnValue( + expectedParsedArgs, + ); + mockPostHogController.capture.mockResolvedValue( + expectedControllerResponse, + ); + + // Act + const result = await captureFunction(inputArgs); + + // Assert + expect( + mockPostHogLlmPropertiesPayloadSchema.parse, + ).toHaveBeenCalledWith(inputArgs); + expect(mockPostHogController.capture).toHaveBeenCalledWith({ + eventName: '$ai_generation', + distinctId: 'test-user-123', + properties: { + userId: 'test-user-123', + $ai_model: 'gpt-4', + $ai_provider: 'openai', + $ai_input: 'test input', + $ai_output_choices: ['choice1', 'choice2'], + $ai_trace_id: 'trace-123', + $ai_input_tokens: 100, + $ai_output_tokens: 200, + $ai_latency: 1500, + $ai_http_status: 200, + $ai_base_url: 'https://api.openai.com', + }, + }); + + expect(result).toEqual({ + content: [ + { + type: 'text', + text: 'Successfully captured LLM observability data', + }, + ], + }); + + expect(mockDebug).toHaveBeenCalledWith( + 'Capture LLM Observability in PostHog...', + inputArgs, + ); + expect(mockDebug).toHaveBeenCalledWith( + 'Got the response from the controller', + expectedControllerResponse, + ); + }); + + it('should handle minimal required fields only', async () => { + // Arrange + const inputArgs: PostHogLlmPropertiesPayloadSchemaType = { + userId: 'test-user-123', + model: 'gpt-4', + provider: 'openai', + }; + + const expectedParsedArgs = { ...inputArgs }; + const expectedControllerResponse = { + content: 'Successfully captured minimal LLM data', + }; + + mockPostHogLlmPropertiesPayloadSchema.parse.mockReturnValue( + expectedParsedArgs, + ); + mockPostHogController.capture.mockResolvedValue( + expectedControllerResponse, + ); + + // Act + const result = await captureFunction(inputArgs); + + // Assert + expect(mockPostHogController.capture).toHaveBeenCalledWith({ + eventName: '$ai_generation', + distinctId: 'test-user-123', + properties: { + userId: 'test-user-123', + $ai_model: 'gpt-4', + $ai_provider: 'openai', + }, + }); + + expect(result).toEqual({ + content: [ + { + type: 'text', + text: 'Successfully captured minimal LLM data', + }, + ], + }); + }); + + it('should handle undefined optional fields correctly', async () => { + // Arrange + const inputArgs: PostHogLlmPropertiesPayloadSchemaType = { + userId: 'test-user-123', + model: 'gpt-4', + provider: 'openai', + input: 'test input', + outputChoices: undefined, + traceId: undefined, + inputTokens: 100, + outputTokens: undefined, + latency: undefined, + httpStatus: undefined, + baseUrl: undefined, + }; + + const expectedParsedArgs = { ...inputArgs }; + const expectedControllerResponse = { + content: 'Successfully captured partial LLM data', + }; + + mockPostHogLlmPropertiesPayloadSchema.parse.mockReturnValue( + expectedParsedArgs, + ); + mockPostHogController.capture.mockResolvedValue( + expectedControllerResponse, + ); + + // Act + const result = await captureFunction(inputArgs); + + // Assert + expect(mockPostHogController.capture).toHaveBeenCalledWith({ + eventName: '$ai_generation', + distinctId: 'test-user-123', + properties: { + userId: 'test-user-123', + $ai_model: 'gpt-4', + $ai_provider: 'openai', + $ai_input: 'test input', + $ai_input_tokens: 100, + }, + }); + + expect(result).toEqual({ + content: [ + { + type: 'text', + text: 'Successfully captured partial LLM data', + }, + ], + }); + }); + + it('should handle schema validation errors', async () => { + // Arrange + const inputArgs: PostHogLlmPropertiesPayloadSchemaType = { + userId: 'test-user-123', + model: 'gpt-4', + provider: 'openai', + }; + + const schemaError = new Error('Schema validation failed'); + const formattedError = { + content: [ + { + type: 'text' as const, + text: 'Schema validation error: Schema validation failed', + }, + ], + }; + + mockPostHogLlmPropertiesPayloadSchema.parse.mockImplementation( + () => { + throw schemaError; + }, + ); + mockFormatErrorForMcpTool.mockReturnValue(formattedError); + + // Act + const result = await captureFunction(inputArgs); + + // Assert + expect( + mockPostHogLlmPropertiesPayloadSchema.parse, + ).toHaveBeenCalledWith(inputArgs); + expect(mockPostHogController.capture).not.toHaveBeenCalled(); + expect(mockFormatErrorForMcpTool).toHaveBeenCalledWith(schemaError); + expect(mockError).toHaveBeenCalledWith( + 'Error tracking LLM generation in PostHog', + schemaError, + ); + expect(result).toEqual(formattedError); + }); + + it('should handle controller errors', async () => { + // Arrange + const inputArgs: PostHogLlmPropertiesPayloadSchemaType = { + userId: 'test-user-123', + model: 'gpt-4', + provider: 'openai', + }; + + const controllerError = new Error('PostHog API error'); + const formattedError = { + content: [ + { + type: 'text' as const, + text: 'Controller error: PostHog API error', + }, + ], + }; + + mockPostHogLlmPropertiesPayloadSchema.parse.mockReturnValue( + inputArgs, + ); + mockPostHogController.capture.mockRejectedValue(controllerError); + mockFormatErrorForMcpTool.mockReturnValue(formattedError); + + // Act + const result = await captureFunction(inputArgs); + + // Assert + expect(mockPostHogController.capture).toHaveBeenCalled(); + expect(mockFormatErrorForMcpTool).toHaveBeenCalledWith( + controllerError, + ); + expect(mockError).toHaveBeenCalledWith( + 'Error tracking LLM generation in PostHog', + controllerError, + ); + expect(result).toEqual(formattedError); + }); + + it('should create logger with correct context', async () => { + // Arrange + const inputArgs: PostHogLlmPropertiesPayloadSchemaType = { + userId: 'test-user-123', + model: 'gpt-4', + provider: 'openai', + }; + + mockPostHogLlmPropertiesPayloadSchema.parse.mockReturnValue( + inputArgs, + ); + mockPostHogController.capture.mockResolvedValue({ + content: 'success', + }); + + // Act + await captureFunction(inputArgs); + + // Assert + expect(Logger.forContext).toHaveBeenCalledWith( + 'tools/posthog-llm.tool.ts', + 'capturePosthogLlmObservability', + ); + }); + + it('should map all possible fields to PostHog properties correctly', async () => { + // Arrange + const inputArgs: PostHogLlmPropertiesPayloadSchemaType = { + userId: 'test-user-123', + model: 'claude-3', + provider: 'anthropic', + input: 'test prompt', + outputChoices: ['response1', 'response2'], + traceId: 'trace-abc-123', + inputTokens: 50, + outputTokens: 150, + latency: 2000, + httpStatus: 201, + baseUrl: 'https://api.anthropic.com', + }; + + mockPostHogLlmPropertiesPayloadSchema.parse.mockReturnValue( + inputArgs, + ); + mockPostHogController.capture.mockResolvedValue({ + content: 'success', + }); + + // Act + await captureFunction(inputArgs); + + // Assert + expect(mockPostHogController.capture).toHaveBeenCalledWith({ + eventName: '$ai_generation', + distinctId: 'test-user-123', + properties: { + userId: 'test-user-123', + $ai_model: 'claude-3', + $ai_provider: 'anthropic', + $ai_input: 'test prompt', + $ai_output_choices: ['response1', 'response2'], + $ai_trace_id: 'trace-abc-123', + $ai_input_tokens: 50, + $ai_output_tokens: 150, + $ai_latency: 2000, + $ai_http_status: 201, + $ai_base_url: 'https://api.anthropic.com', + }, + }); + }); + }); +}); diff --git a/src/tools/posthog-llm.tool.ts b/src/tools/posthog-llm.tool.ts index 3cc6968..43589f3 100644 --- a/src/tools/posthog-llm.tool.ts +++ b/src/tools/posthog-llm.tool.ts @@ -53,12 +53,10 @@ async function capturePosthogLlmObservability( keyof PostHogLlmPropertiesPayloadSchemaType >) { if (trackArgs[key] !== undefined) { - const posthogKey = toPostHogKey[key]; - if (posthogKey) { - // Type assertion to satisfy TS: posthogKey is a key of posthogProperties - (posthogProperties as Record)[posthogKey] = - trackArgs[key]; - } + const posthogKey = toPostHogKey[key]!; + // Type assertion to satisfy TS: posthogKey is a key of posthogProperties + (posthogProperties as Record)[posthogKey] = + trackArgs[key]; } } diff --git a/src/utils/__mocks__/logger.util.ts b/src/utils/__mocks__/logger.util.ts new file mode 100644 index 0000000..6506a3a --- /dev/null +++ b/src/utils/__mocks__/logger.util.ts @@ -0,0 +1,18 @@ +export const mockDebug = jest.fn(); +export const mockError = jest.fn(); +export const mockInfo = jest.fn(); +export const mockWarn = jest.fn(); + +const mockLogger = { + debug: mockDebug, + error: mockError, + info: mockInfo, + warn: mockWarn, + forMethod: jest.fn().mockReturnThis(), +}; + +export const Logger = { + forContext: jest.fn(() => mockLogger), + getSessionId: jest.fn(() => 'mock-session-id'), + getLogFilePath: jest.fn(() => '/mock/log/path'), +}; diff --git a/src/utils/logger.util.ts b/src/utils/logger.util.ts index 851aa89..ef0274a 100644 --- a/src/utils/logger.util.ts +++ b/src/utils/logger.util.ts @@ -311,9 +311,9 @@ class Logger { } if (process.env.NODE_ENV === 'test') { - console[level](logMessage); + console[level]?.(logMessage); } else { - console.error(logMessage); + console.error?.(logMessage); } }