andy-stack_vaultkeeper-ai/__tests__/Services/OrchestrationAgent.test.ts

import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { OrchestrationAgent } from '../../Services/AIServices/OrchestrationAgent';
import { RegisterSingleton, DeregisterAllServices } from '../../Services/DependencyService';
import { Services } from '../../Services/Services';
import { AITool } from '../../Enums/AITool';
import { AIToolResponse } from '../../AIClasses/ToolDefinitions/AIToolResponse';
import { AIToolResponsePayload } from 'AIClasses/ToolDefinitions/AIToolResponsePayload';

/**
 * UNIT TESTS - OrchestrationAgent
 *
 * Tests the orchestration agent's basic functionality:
 * - Service initialization
 * - Agent configuration
 *
 * Note: Full workflow integration tests with nested agents are in MultiAgentIntegration.test.ts
 * These tests focus on the OrchestrationAgent's setup and configuration.
 */

describe('OrchestrationAgent - Unit Tests', () => {
	let service: OrchestrationAgent;
	let mockAI: any;
	let mockPrompt: any;
	let mockAIToolService: any;

	beforeEach(() => {
		// Mock IPrompt
		mockPrompt = {
			systemInstruction: vi.fn().mockReturnValue('System instruction'),
			userInstruction: vi.fn().mockResolvedValue('User instruction'),
			planningInstruction: vi.fn().mockResolvedValue('Planning instruction'),
			orchestrationInstruction: vi.fn().mockResolvedValue('Orchestration instruction'),
			executionInstruction: vi.fn().mockReturnValue('Execution instruction')
		};
		RegisterSingleton(Services.IPrompt, mockPrompt);

		// Mock AIToolService
		mockAIToolService = {
			performAITool: vi.fn().mockResolvedValue(
				new AIToolResponse(AITool.SearchVaultFiles, new AIToolResponsePayload({ results: [] }), 'test-tool-id')
			)
		};
		RegisterSingleton(Services.AIToolService, mockAIToolService);

		// Mock IAIClass
		mockAI = {
			systemPrompt: '',
			userInstruction: '',
			toolDefinitions: [],
			agentType: undefined,
			streamRequest: vi.fn()
		};
		RegisterSingleton(Services.IAIClass, mockAI);

		// Mock SettingsService
		RegisterSingleton(Services.SettingsService, {
			settings: { enableMemories: false, allowUpdatingMemories: false }
		});

		// Create service
		service = new OrchestrationAgent();
	});

	afterEach(() => {
		DeregisterAllServices();
		vi.restoreAllMocks();
	});

	describe('constructor', () => {
		it('should initialize with all required services', () => {
			const testService = new OrchestrationAgent();

			expect(testService).toBeDefined();
		});
	});

	describe('resolveAIProvider', () => {
		it('should resolve AI provider successfully', () => {
			service.resolveAIProvider();

			expect(service).toBeDefined();
		});

		it('should set up AI provider from dependency service', () => {
			service.resolveAIProvider();

			// Verify the service can be used
			expect(service).toHaveProperty('runPlannedWorkflow');
		});
	});

	describe('service methods', () => {
		it('should have runPlannedWorkflow method', () => {
			expect(service.runPlannedWorkflow).toBeDefined();
			expect(typeof service.runPlannedWorkflow).toBe('function');
		});

		it('should throw error if AI provider not resolved before running workflow', async () => {
			const callbacks = {
				onSubmit: vi.fn(),
				onStreamingUpdate: vi.fn(),
				onThoughtUpdate: vi.fn(),
				onPlanningStarted: vi.fn(),
				onPlanningFinished: vi.fn(),
				onUserQuestion: vi.fn(),
				onPlanUpdate: vi.fn(),
				onPlanStepUpdate: vi.fn(),
				onPlanReset: vi.fn(),
				onComplete: vi.fn(),
				onCancel: vi.fn()
			};

			// Note: runPlannedWorkflow spawns new agents internally which get
			// their own AI provider from the dependency service, so this test
			// verifies the service is properly constructed
			expect(service.runPlannedWorkflow).toBeDefined();
			expect(typeof service.runPlannedWorkflow).toBe('function');
		});
	});
});