n8n/packages/@n8n/ai-workflow-builder.ee/test/test-utils.ts

import type { ToolRunnableConfig } from '@langchain/core/tools';
import type { LangGraphRunnableConfig } from '@langchain/langgraph';
import { getCurrentTaskInput } from '@langchain/langgraph';
import type { MockProxy } from 'jest-mock-extended';
import { mock } from 'jest-mock-extended';
import type {
	INode,
	INodeTypeDescription,
	INodeParameters,
	IConnection,
	NodeConnectionType,
	INodeListSearchResult,
	IRunData,
	ITaskDataConnections,
	NodeExecutionSchema,
	Schema,
	IDataObject,
} from 'n8n-workflow';
import { jsonParse } from 'n8n-workflow';

import type { ProgrammaticEvaluationResult } from '@/validation/types';

import type { ResourceLocatorCallback } from '../src/types/callbacks';
import type { ProgressReporter, ToolProgressMessage } from '../src/types/tools';
import type { SimpleWorkflow } from '../src/types/workflow';

export const mockProgress = (): MockProxy<ProgressReporter> => mock<ProgressReporter>();

// Mock state helpers
export const mockStateHelpers = () => ({
	getNodes: jest.fn(() => [] as INode[]),
	getConnections: jest.fn(() => ({}) as SimpleWorkflow['connections']),
	updateNode: jest.fn((_id: string, _updates: Partial<INode>) => undefined),
	addNodes: jest.fn((_nodes: INode[]) => undefined),
	removeNode: jest.fn((_id: string) => undefined),
	addConnections: jest.fn((_connections: IConnection[]) => undefined),
	removeConnection: jest.fn((_sourceId: string, _targetId: string, _type?: string) => undefined),
});

export type MockStateHelpers = ReturnType<typeof mockStateHelpers>;

// Simple node creation helper
export const createNode = (overrides: Partial<INode> = {}): INode => ({
	id: 'node1',
	name: 'TestNode',
	type: 'n8n-nodes-base.code',
	typeVersion: 1,
	position: [0, 0],
	...overrides,
	// Ensure parameters are properly merged if provided in overrides
	parameters: overrides.parameters ?? {},
});

// Simple workflow builder
export const createWorkflow = (nodes: INode[] = []): SimpleWorkflow => {
	const workflow: SimpleWorkflow = { nodes, connections: {}, name: 'Test workflow' };
	return workflow;
};

// Create mock node type description
export const createNodeType = (
	overrides: Partial<INodeTypeDescription> = {},
): INodeTypeDescription => ({
	displayName: overrides.displayName ?? 'Test Node',
	name: overrides.name ?? 'test.node',
	group: overrides.group ?? ['transform'],
	version: overrides.version ?? 1,
	description: overrides.description ?? 'Test node description',
	defaults: overrides.defaults ?? { name: 'Test Node' },
	inputs: overrides.inputs ?? ['main'],
	outputs: overrides.outputs ?? ['main'],
	properties: overrides.properties ?? [],
	...overrides,
});

// Common node types for testing
export const nodeTypes = {
	code: createNodeType({
		displayName: 'Code',
		name: 'n8n-nodes-base.code',
		group: ['transform'],
		properties: [
			{
				displayName: 'JavaScript',
				name: 'jsCode',
				type: 'string',
				typeOptions: {
					editor: 'codeNodeEditor',
				},
				default: '',
			},
		],
	}),
	httpRequest: createNodeType({
		displayName: 'HTTP Request',
		name: 'n8n-nodes-base.httpRequest',
		group: ['input'],
		properties: [
			{
				displayName: 'URL',
				name: 'url',
				type: 'string',
				default: '',
			},
			{
				displayName: 'Method',
				name: 'method',
				type: 'options',
				options: [
					{ name: 'GET', value: 'GET' },
					{ name: 'POST', value: 'POST' },
				],
				default: 'GET',
			},
		],
	}),
	webhook: createNodeType({
		displayName: 'Webhook',
		name: 'n8n-nodes-base.webhook',
		group: ['trigger'],
		inputs: [],
		outputs: ['main'],
		webhooks: [
			{
				name: 'default',
				httpMethod: 'POST',
				responseMode: 'onReceived',
				path: 'webhook',
			},
		],
		properties: [
			{
				displayName: 'Path',
				name: 'path',
				type: 'string',
				default: 'webhook',
			},
		],
	}),
	agent: createNodeType({
		displayName: 'AI Agent',
		name: '@n8n/n8n-nodes-langchain.agent',
		group: ['output'],
		inputs: ['ai_agent'],
		outputs: ['main'],
		properties: [],
	}),
	openAiModel: createNodeType({
		displayName: 'OpenAI Chat Model',
		name: '@n8n/n8n-nodes-langchain.lmChatOpenAi',
		group: ['output'],
		inputs: [],
		outputs: ['ai_languageModel'],
		properties: [],
	}),
	setNode: createNodeType({
		displayName: 'Set',
		name: 'n8n-nodes-base.set',
		group: ['transform'],
		properties: [
			{
				displayName: 'Values to Set',
				name: 'values',
				type: 'collection',
				default: {},
			},
		],
	}),
	ifNode: createNodeType({
		displayName: 'If',
		name: 'n8n-nodes-base.if',
		group: ['transform'],
		inputs: ['main'],
		outputs: ['main', 'main'],
		outputNames: ['true', 'false'],
		properties: [
			{
				displayName: 'Conditions',
				name: 'conditions',
				type: 'collection',
				default: {},
			},
		],
	}),
	mergeNode: createNodeType({
		displayName: 'Merge',
		name: 'n8n-nodes-base.merge',
		group: ['transform'],
		inputs: ['main', 'main'],
		outputs: ['main'],
		inputNames: ['Input 1', 'Input 2'],
		properties: [
			{
				displayName: 'Mode',
				name: 'mode',
				type: 'options',
				options: [
					{ name: 'Append', value: 'append' },
					{ name: 'Merge By Index', value: 'mergeByIndex' },
					{ name: 'Merge By Key', value: 'mergeByKey' },
				],
				default: 'append',
			},
		],
	}),
	vectorStoreNode: createNodeType({
		displayName: 'Vector Store',
		name: '@n8n/n8n-nodes-langchain.vectorStore',
		subtitle: '={{$parameter["mode"] === "retrieve" ? "Retrieve" : "Insert"}}',
		group: ['transform'],
		inputs: `={{ ((parameter) => {
			function getInputs(parameters) {
				const mode = parameters?.mode;
				const inputs = [];
				if (mode === 'retrieve-as-tool') {
					inputs.push({
						displayName: 'Embedding',
						type: 'ai_embedding',
						required: true
					});
				} else {
					inputs.push({
						displayName: '',
						type: 'main'
					});
					inputs.push({
						displayName: 'Embedding',
						type: 'ai_embedding',
						required: true
					});
				}
				return inputs;
			};
			return getInputs(parameter)
		})($parameter) }}`,
		outputs: `={{ ((parameter) => {
			function getOutputs(parameters) {
				const mode = parameters?.mode;
				if (mode === 'retrieve-as-tool') {
					return ['ai_tool'];
				} else if (mode === 'retrieve') {
					return ['ai_document'];
				} else {
					return ['main'];
				}
			};
			return getOutputs(parameter)
		})($parameter) }}`,
		properties: [
			{
				displayName: 'Mode',
				name: 'mode',
				type: 'options',
				options: [
					{ name: 'Insert', value: 'insert' },
					{ name: 'Retrieve', value: 'retrieve' },
					{ name: 'Retrieve (As Tool)', value: 'retrieve-as-tool' },
				],
				default: 'insert',
			},
			// Many more properties would be here in reality
		],
	}),
};

// Helper to create connections
export const createConnection = (
	_fromId: string,
	toId: string,
	type: NodeConnectionType = 'main',
	index: number = 0,
) => ({
	node: toId,
	type,
	index,
});

// Generic chain interface
interface Chain<TInput = Record<string, unknown>, TOutput = Record<string, unknown>> {
	invoke: (input: TInput) => Promise<TOutput>;
}

// Generic mock chain factory with proper typing
export const mockChain = <
	TInput = Record<string, unknown>,
	TOutput = Record<string, unknown>,
>(): MockProxy<Chain<TInput, TOutput>> => {
	return mock<Chain<TInput, TOutput>>();
};

// Convenience factory for parameter updater chain
export const mockParameterUpdaterChain = () => {
	return mockChain<Record<string, unknown>, { parameters: Record<string, unknown> }>();
};

// Helper to assert node parameters
export const expectNodeToHaveParameters = (
	node: INode,
	expectedParams: Partial<INodeParameters>,
): void => {
	expect(node.parameters).toMatchObject(expectedParams);
};

// Helper to assert connections exist
export const expectConnectionToExist = (
	connections: SimpleWorkflow['connections'],
	fromId: string,
	toId: string,
	type: string = 'main',
): void => {
	expect(connections[fromId]).toBeDefined();
	expect(connections[fromId][type]).toBeDefined();
	expect(connections[fromId][type]).toContainEqual(
		expect.arrayContaining([expect.objectContaining({ node: toId })]),
	);
};

// ========== LangGraph Testing Utilities ==========

// Types for mocked Command results
export type MockedCommandResult = { content: string };

// Common parsed content structure for tool results
export interface ParsedToolContent {
	update: {
		messages: Array<{ kwargs: { content: string } }>;
		workflowOperations?: Array<{
			type: string;
			nodes?: INode[];
			[key: string]: unknown;
		}>;
		workflowValidation?: ProgrammaticEvaluationResult | null;
	};
}

// Setup LangGraph mocks
export const setupLangGraphMocks = () => {
	const mockGetCurrentTaskInput = getCurrentTaskInput as jest.MockedFunction<
		typeof getCurrentTaskInput
	>;

	jest.mock('@langchain/langgraph', () => ({
		getCurrentTaskInput: jest.fn(),
		Command: jest.fn().mockImplementation((params: Record<string, unknown>) => ({
			content: JSON.stringify(params),
		})),
	}));

	return { mockGetCurrentTaskInput };
};

// Parse tool result with double-wrapped content handling
export const parseToolResult = <T = ParsedToolContent>(result: unknown): T => {
	const parsed = jsonParse<{ content?: string }>((result as MockedCommandResult).content);
	return parsed.content ? jsonParse<T>(parsed.content) : (parsed as T);
};

// ========== Progress Message Utilities ==========

// Extract progress messages from mockWriter
export const extractProgressMessages = (
	mockWriter: jest.Mock,
): Array<ToolProgressMessage<string>> => {
	const progressCalls: Array<ToolProgressMessage<string>> = [];

	mockWriter.mock.calls.forEach((call) => {
		// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
		const [arg] = call;
		progressCalls.push(arg as ToolProgressMessage<string>);
	});
	return progressCalls;
};

// Find specific progress message by type
export const findProgressMessage = (
	messages: Array<ToolProgressMessage<string>>,
	status: 'running' | 'completed' | 'error',
	updateType?: string,
): ToolProgressMessage<string> | undefined => {
	return messages.find(
		(msg) => msg.status === status && (!updateType || msg.updates[0]?.type === updateType),
	);
};

// ========== Tool Config Helpers ==========

// Create basic tool config
export const createToolConfig = (
	toolName: string,
	callId: string = 'test-call',
): ToolRunnableConfig => ({
	toolCall: { id: callId, name: toolName, args: {} },
});

// Create tool config with writer for progress tracking
export const createToolConfigWithWriter = (
	toolName: string,
	callId: string = 'test-call',
): ToolRunnableConfig & LangGraphRunnableConfig & { writer: jest.Mock } => {
	const mockWriter = jest.fn();
	return {
		toolCall: { id: callId, name: toolName, args: {} },
		writer: mockWriter,
	};
};

// ========== Workflow State Helpers ==========

// Setup workflow state with mockGetCurrentTaskInput
export const setupWorkflowState = (
	mockGetCurrentTaskInput: jest.MockedFunction<typeof getCurrentTaskInput>,
	workflow: SimpleWorkflow = createWorkflow([]),
) => {
	mockGetCurrentTaskInput.mockReturnValue({
		workflowJSON: workflow,
		workflowOperations: null,
		workflowContext: {},
		workflowValidation: null,
		messages: [],
		previousSummary: 'EMPTY',
	});
};

// ========== Execution Data Builders ==========

// Build run data entry from simple JSON data
export interface MockRunDataEntry {
	json: Record<string, unknown>;
	startTime?: number;
	executionTime?: number;
}

// Create mock run data from simplified entries
export const createMockRunData = (entries: Record<string, MockRunDataEntry[]>): IRunData => {
	const runData: IRunData = {};
	let executionIndex = 0;
	for (const [nodeName, items] of Object.entries(entries)) {
		runData[nodeName] = [
			{
				data: {
					main: [items.map((item) => ({ json: item.json }))] as ITaskDataConnections['main'],
				},
				startTime: items[0]?.startTime ?? Date.now(),
				executionTime: items[0]?.executionTime ?? 100,
				executionIndex: executionIndex++,
				source: [null],
			},
		];
	}
	return runData;
};

// Create mock execution schema from simplified entries
export interface MockNodeSchema {
	nodeName: string;
	schema: Schema;
}

export const createMockExecutionSchema = (nodeSchemas: MockNodeSchema[]): NodeExecutionSchema[] => {
	return nodeSchemas.map(({ nodeName, schema }) => ({
		nodeName,
		schema,
	}));
};

// Helper to create a Schema object for testing
export const createMockSchema = (
	type: Schema['type'],
	path: string,
	value: Schema['value'],
	key?: string,
): Schema => ({
	type,
	path,
	value,
	...(key && { key }),
});

// Generate large test data for truncation tests
export const createLargeTestData = (itemCount = 100, fieldValueSize = 30): IDataObject[] => {
	return Array.from({ length: itemCount }, (_, i) => ({
		id: i,
		field: 'x'.repeat(fieldValueSize) + String(i),
		extra: 'y'.repeat(fieldValueSize),
	}));
};

// ========== Extended Workflow State Setup ==========

export interface ExecutionDataOptions {
	runData?: IRunData;
	lastNodeExecuted?: string;
	error?: { message: string; description?: string };
}

export interface ExpressionValueTestData {
	expression: string;
	resolvedValue: unknown;
	nodeType?: string;
}

export interface WorkflowStateOptions {
	workflow: SimpleWorkflow;
	executionData?: ExecutionDataOptions;
	executionSchema?: NodeExecutionSchema[];
	expressionValues?: Record<string, ExpressionValueTestData[]>;
}

// Setup workflow state with execution context (extended version)
export const setupWorkflowStateWithContext = (
	mockGetCurrentTaskInput: jest.MockedFunction<typeof getCurrentTaskInput>,
	options: WorkflowStateOptions,
) => {
	mockGetCurrentTaskInput.mockReturnValue({
		workflowJSON: options.workflow,
		workflowOperations: null,
		workflowContext: {
			executionData: options.executionData ?? null,
			executionSchema: options.executionSchema ?? null,
			expressionValues: options.expressionValues ?? null,
		},
		workflowValidation: null,
		messages: [],
		previousSummary: 'EMPTY',
	});
};

// ========== AI Workflow Helpers ==========

// Setup AI connections on workflow (e.g., model -> agent)
export const setupAIWorkflowConnections = (
	workflow: SimpleWorkflow,
	modelNodeName: string,
	agentNodeName: string,
	connectionType: NodeConnectionType = 'ai_languageModel',
) => {
	workflow.connections[modelNodeName] = {
		[connectionType]: [[{ node: agentNodeName, type: connectionType, index: 0 }]],
	};
};

// ========== Common Tool Assertions ==========

// Expect tool success message
export const expectToolSuccess = (
	content: ParsedToolContent,
	expectedMessage: string | RegExp,
): void => {
	const message = content.update.messages[0]?.kwargs.content;
	expect(message).toBeDefined();
	if (typeof expectedMessage === 'string') {
		expect(message).toContain(expectedMessage);
	} else {
		expect(message).toMatch(expectedMessage);
	}
};

// Expect tool error message
export const expectToolError = (
	content: ParsedToolContent,
	expectedError: string | RegExp,
): void => {
	const message = content.update.messages[0]?.kwargs.content;
	if (typeof expectedError === 'string') {
		expect(message).toBe(expectedError);
	} else {
		expect(message).toMatch(expectedError);
	}
};

// Expect workflow operation of specific type
export const expectWorkflowOperation = (
	content: ParsedToolContent,
	operationType: string,
	matcher?: Record<string, unknown>,
): void => {
	const operation = content.update.workflowOperations?.[0];
	expect(operation).toBeDefined();
	expect(operation?.type).toBe(operationType);
	if (matcher) {
		expect(operation).toMatchObject(matcher);
	}
};

// Expect node was added
export const expectNodeAdded = (content: ParsedToolContent, expectedNode: Partial<INode>): void => {
	expectWorkflowOperation(content, 'addNodes');
	const addedNode = content.update.workflowOperations?.[0]?.nodes?.[0];
	expect(addedNode).toBeDefined();
	expect(addedNode).toMatchObject(expectedNode);
};

// Expect node was removed
export const expectNodeRemoved = (content: ParsedToolContent, nodeId: string): void => {
	expectWorkflowOperation(content, 'removeNode', { nodeIds: [nodeId] });
};

// Expect connections were added
export const expectConnectionsAdded = (
	content: ParsedToolContent,
	expectedCount?: number,
): void => {
	expectWorkflowOperation(content, 'addConnections');
	if (expectedCount !== undefined) {
		const connections = content.update.workflowOperations?.[0]?.connections;
		expect(connections).toHaveLength(expectedCount);
	}
};

// Expect node was updated
export const expectNodeUpdated = (
	content: ParsedToolContent,
	nodeId: string,
	expectedUpdates?: Record<string, unknown>,
): void => {
	expectWorkflowOperation(content, 'updateNode', {
		nodeId,
		...(expectedUpdates ? { updates: expect.objectContaining(expectedUpdates) } : {}),
	});
};

// ========== Test Data Builders ==========

// Build add node input
export const buildAddNodeInput = (overrides: {
	nodeType: string;
	nodeVersion?: number;
	name?: string;
	initialParametersReasoning?: string;
	initialParameters?: Record<string, unknown>;
}) => ({
	nodeType: overrides.nodeType,
	nodeVersion: overrides.nodeVersion ?? 1,
	name: overrides.name ?? 'Test Node',
	initialParametersReasoning:
		overrides.initialParametersReasoning ??
		'Standard node with static inputs/outputs, no initial parameters needed',
	initialParameters: overrides.initialParameters ?? {},
});

// Build connect nodes input
export const buildConnectNodesInput = (overrides: {
	sourceNodeId: string;
	targetNodeId: string;
	sourceOutputIndex?: number;
	targetInputIndex?: number;
}) => ({
	sourceNodeId: overrides.sourceNodeId,
	targetNodeId: overrides.targetNodeId,
	sourceOutputIndex: overrides.sourceOutputIndex ?? 0,
	targetInputIndex: overrides.targetInputIndex ?? 0,
});

// Build node search query
export const buildNodeSearchQuery = (
	queryType: 'name' | 'subNodeSearch',
	query?: string,
	connectionType?: NodeConnectionType,
) => ({
	queryType,
	...(query && { query }),
	...(connectionType && { connectionType }),
});

// Build update node parameters input
export const buildUpdateNodeInput = (nodeId: string, changes: string[]) => ({
	nodeId,
	changes,
});

// Build node details input
export const buildNodeDetailsInput = (overrides: {
	nodeName: string;
	nodeVersion?: number;
	withParameters?: boolean;
	withConnections?: boolean;
}) => ({
	nodeName: overrides.nodeName,
	nodeVersion: overrides.nodeVersion ?? 1,
	withParameters: overrides.withParameters ?? false,
	withConnections: overrides.withConnections ?? true,
});

// Expect node details in response
export const expectNodeDetails = (
	content: ParsedToolContent,
	expectedDetails: Partial<{
		name: string;
		displayName: string;
		description: string;
		subtitle?: string;
	}>,
): void => {
	const message = content.update.messages[0]?.kwargs.content;
	expect(message).toBeDefined();

	// Check for expected XML-like tags in formatted output
	if (expectedDetails.name) {
		expect(message).toContain(`<name>${expectedDetails.name}</name>`);
	}
	if (expectedDetails.displayName) {
		expect(message).toContain(`<display_name>${expectedDetails.displayName}</display_name>`);
	}
	if (expectedDetails.description) {
		expect(message).toContain(`<description>${expectedDetails.description}</description>`);
	}
	if (expectedDetails.subtitle) {
		expect(message).toContain(`<subtitle>${expectedDetails.subtitle}</subtitle>`);
	}
};

// Helper to validate XML-like structure in output
export const expectXMLTag = (
	content: string,
	tagName: string,
	expectedValue?: string | RegExp,
): void => {
	const tagRegex = new RegExp(`<${tagName}>([\\s\\S]*?)</${tagName}>`);
	const match = content.match(tagRegex);
	expect(match).toBeDefined();
	if (expectedValue) {
		if (typeof expectedValue === 'string') {
			expect(match?.[1]?.trim()).toBe(expectedValue);
		} else {
			expect(match?.[1]).toMatch(expectedValue);
		}
	}
};

// Common reasoning strings
export const REASONING = {
	STATIC_NODE: 'Node has static inputs/outputs, no connection parameters needed',
	DYNAMIC_AI_NODE: 'AI node has dynamic inputs, setting connection parameters',
	TRIGGER_NODE: 'Trigger node, no connection parameters needed',
	WEBHOOK_NODE: 'Webhook is a trigger node, no connection parameters needed',
} as const;

// ========== Resource Locator Testing Utilities ==========

// Create a node type with resource locator property
export const createNodeTypeWithResourceLocator = (
	nodeName: string,
	parameterName: string,
	searchMethod: string,
	overrides: Partial<INodeTypeDescription> = {},
): INodeTypeDescription =>
	createNodeType({
		name: nodeName,
		displayName: overrides.displayName ?? 'Test Resource Node',
		credentials: overrides.credentials ?? [{ name: 'testApi', required: true }],
		properties: [
			{
				displayName: 'Resource',
				name: parameterName,
				type: 'resourceLocator',
				default: { mode: 'list', value: '' },
				modes: [
					{
						displayName: 'From List',
						name: 'list',
						type: 'list',
						typeOptions: {
							searchListMethod: searchMethod,
							searchable: true,
						},
					},
					{
						displayName: 'By ID',
						name: 'id',
						type: 'string',
					},
				],
			},
			...(overrides.properties ?? []),
		],
		...overrides,
	});

// Create a mock resource locator callback
export const mockResourceLocatorCallback = (
	results: INodeListSearchResult = { results: [] },
): jest.MockedFunction<ResourceLocatorCallback> => {
	return jest.fn().mockResolvedValue(results);
};

// Create sample resource locator results
export const createResourceLocatorResults = (
	items: Array<{ name: string; value: string; description?: string }>,
	paginationToken?: string,
): INodeListSearchResult => ({
	results: items.map((item) => ({
		name: item.name,
		value: item.value,
		...(item.description && { description: item.description }),
	})),
	...(paginationToken && { paginationToken }),
});