Added unit tests for posthog-llm.tool.ts
This commit is contained in:
433
src/tools/posthog-llm.tool.test.ts
Normal file
433
src/tools/posthog-llm.tool.test.ts
Normal 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',
|
||||
},
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -53,14 +53,12 @@ async function capturePosthogLlmObservability(
|
||||
keyof PostHogLlmPropertiesPayloadSchemaType
|
||||
>) {
|
||||
if (trackArgs[key] !== undefined) {
|
||||
const posthogKey = toPostHogKey[key];
|
||||
if (posthogKey) {
|
||||
const posthogKey = toPostHogKey[key]!;
|
||||
// Type assertion to satisfy TS: posthogKey is a key of posthogProperties
|
||||
(posthogProperties as Record<string, unknown>)[posthogKey] =
|
||||
trackArgs[key];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// validated input for the controller
|
||||
const posthogInput: PostHogLlmInputSchemaType = {
|
||||
|
||||
18
src/utils/__mocks__/logger.util.ts
Normal file
18
src/utils/__mocks__/logger.util.ts
Normal 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'),
|
||||
};
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user