Home/llm-observability-mcp
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added unit tests for posthog-llm.tool.ts

Browse Source
This commit is contained in:
Stefano Fiorini
2025-07-17 11:35:22 -05:00
parent 7ff3aa7991
commit 72d55a5b40
4 changed files with 457 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

433
src/tools/posthog-llm.tool.test.ts Normal file
View File

@@ -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<typeof PostHogController>;
let mockFormatErrorForMcpTool: jest.MockedFunction<
typeof formatErrorForMcpTool
>;
let mockPostHogLlmPropertiesPayloadSchema: jest.Mocked<any>;
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<CallToolResult>;
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',
},
});
});
});
});

10
src/tools/posthog-llm.tool.ts
View File

@@ -53,12 +53,10 @@ async function capturePosthogLlmObservability(
keyof PostHogLlmPropertiesPayloadSchemaType keyof PostHogLlmPropertiesPayloadSchemaType
>) { >) {
if (trackArgs[key] !== undefined) { if (trackArgs[key] !== undefined) {
const posthogKey = toPostHogKey[key]; const posthogKey = toPostHogKey[key]!;
if (posthogKey) { // Type assertion to satisfy TS: posthogKey is a key of posthogProperties
// Type assertion to satisfy TS: posthogKey is a key of posthogProperties (posthogProperties as Record<string, unknown>)[posthogKey] =
(posthogProperties as Record<string, unknown>)[posthogKey] = trackArgs[key];
trackArgs[key];
}
} }
} }

18
src/utils/__mocks__/logger.util.ts Normal file
View File

@@ -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'),
};

4
src/utils/logger.util.ts
View File

@@ -311,9 +311,9 @@ class Logger {
} }
if (process.env.NODE_ENV === 'test') { if (process.env.NODE_ENV === 'test') {
console[level](logMessage); console[level]?.(logMessage);
} else { } else {
console.error(logMessage); console.error?.(logMessage);
} }
} }