cmte/src/core/execution/__tests__/set-executor.test.js

Design by Committee™ except it's just you and LLMs

// Remove hoisted mock
// vi.mock('../../task/template-renderer.js', () => { ... });

import { SetExecutor } from '../set-executor.js';
import { vi, describe, it, expect, beforeEach, test, afterEach } from 'vitest';
import { WorkflowExecutor } from '../workflow-executor.js';
import { ComponentRegistry as Registry } from '../../components/registry.js';
// Import the REAL TemplateRenderer
import { TemplateRenderer } from '../../task/template-renderer.js'; 
// Remove import of mocked version
// import { TemplateRenderer as MockedTemplateRenderer, __mockRender as mockRender } from '../../task/template-renderer.js';

describe('SetExecutor', () => {
  let executor;
  let mockRegistry;
  let mockLLMClient;
  let mockWorkflowExecutor;
  let mockFileManager;
  let mockOutputResolver;

  beforeEach(() => {
    // Remove mockRender clearing and default implementation
    // mockRender.mockClear();
    // mockRender.mockImplementation(...);

    // --- Reset other mocks --- 
    mockRegistry = {
      loadTask: vi.fn(),
      loadSet: vi.fn()
    };
    mockLLMClient = {
      completeMessages: vi.fn()
    };
    mockFileManager = {
        getFiles: vi.fn().mockReturnValue([]),
        readFileRelative: vi.fn().mockResolvedValue(''),
        hasCollection: vi.fn().mockReturnValue(true)
    };
    mockOutputResolver = {
      resolveReference: vi.fn((ref) => {
        throw new Error(`Default Mock resolver: Cannot resolve ${ref}`);
      }),
      setIterationContext: vi.fn(),
      clearIterationContext: vi.fn(),
      registerOutput: vi.fn(),
      registerIteratedOutput: vi.fn()
    };

    mockWorkflowExecutor = {
        registry: mockRegistry,
        workflowPath: '/test/workflow',
        rawOutputPath: '/test/output',
        state: {},
        savePrompts: false,
        dryRun: false,
        apiDryRun: false,
        lite: false,
        useLocalLLM: false,
        mockTaskExecution: false,
        modelConfig: { model: 'default-model-placeholder' },
        llmClient: mockLLMClient,
        fileCollectionManager: mockFileManager,
        outputReferenceResolver: mockOutputResolver,
        recordTaskOutput: vi.fn(),
        isDryRun: vi.fn().mockReturnValue(false),
        addDryRunIssue: vi.fn(),
        isMultiRun: false
    };

    const mockSetDefinition = {
        name: 'test-set',
        tasks: []
    };
    const initialContext = {};

    executor = new SetExecutor(
        mockSetDefinition, 
        mockRegistry, 
        mockWorkflowExecutor, 
        initialContext
    );
  });

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

  describe('task execution', () => {
    // Test: executes multiple tasks
    test('executes multiple tasks', async () => {
      const setDefinition = {
        name: 'test-set',
        tasks: [
          { useTask: 'task1' },
          { useTask: 'task2' }
        ]
      };

      executor.setDefinition = setDefinition;

      mockRegistry.loadTask.mockResolvedValue({
        task: { name: 'mockTask' },
        content: 'Static test content'
      });

      mockLLMClient.completeMessages.mockImplementation(async (messages) => {
        const inputContent = messages[0]?.content || '';
        const mockXmlInput = `<prompt>${inputContent}</prompt>`;
        return `${mockXmlInput}_result`;
      });

      const result = await executor.executeOnce();
      
      expect(mockRegistry.loadTask).toHaveBeenCalledWith('task1');
      expect(mockRegistry.loadTask).toHaveBeenCalledWith('task2');
      expect(mockLLMClient.completeMessages).toHaveBeenCalledTimes(2);
      // Cannot assert on mockRender anymore
      // expect(mockRender).toHaveBeenCalledTimes(2); 
      expect(result.outputs['task1']).toContain('_result');
      expect(result.outputs['task2']).toContain('_result');
      expect(mockWorkflowExecutor.recordTaskOutput).toHaveBeenCalledTimes(2);
      expect(mockWorkflowExecutor.recordTaskOutput).toHaveBeenCalledWith(-1, 'test-set', 'task1', null, expect.stringContaining('<prompt>Static test content</prompt>_result'));
      expect(mockWorkflowExecutor.recordTaskOutput).toHaveBeenCalledWith(-1, 'test-set', 'task2', null, expect.stringContaining('<prompt>Static test content</prompt>_result'));
    });

    // Test: merges prior_outputs
    it('merges prior_outputs with context before rendering', async () => {
      const setDef = {
        name: 'test-set',
        tasks: [{
          useTask: 'task1',
          prior_outputs: {
            priorOutputVar: '{{ someSet.someTask.output }}'
          }
        }]
      };
      const initialContext = { baseVar: 'baseValue' };
      const taskDef = { name: 'task1', content: 'Base: {{baseVar}}, Prior: {{priorOutputVar}}' };
      
      mockRegistry.loadSet.mockResolvedValue(setDef);
      mockRegistry.loadTask.mockResolvedValue({ task: taskDef, content: taskDef.content }); 
      
      mockOutputResolver.resolveReference.mockImplementation((ref) => {
          if (ref === 'someSet.someTask.output') {
              return 'ResolvedPriorOutputValue';
          }
          // Simulate resolver failure for non-references
          throw new Error(`Resolver Mock: '${ref}' is not a valid reference format.`);
      });

      // Update the executor instance with the correct set definition for this test
      executor.setDefinition = setDef;
            
      await executor.executeOnce(initialContext);

      // Cannot assert on mockRender anymore
      // expect(mockRender).toHaveBeenCalledTimes(1);
      // const renderCallArgs = mockRender.mock.calls[0];
      // const templateArg = renderCallArgs[0]; 
      // const contextArg = renderCallArgs[1]; 
      // expect(templateArg).toBe(...);
      // expect(contextArg).toEqual(...);
      
      // Check the final output passed to LLM (which relies on the real render)
      expect(mockLLMClient.completeMessages).toHaveBeenCalledTimes(1);
      const mockLLMCall = mockLLMClient.completeMessages.mock.calls[0];
      const messagesArg = mockLLMCall[0];
      const actualRenderedContent = messagesArg[0].content;

      // Check if the real renderer (with our fix) produced the correct string
      // Note: This assumes LLMXML conversion doesn't drastically change this simple string
      expect(actualRenderedContent).toContain('Base: baseValue, Prior: ResolvedPriorOutputValue');
    });
  });

  describe('validation', () => {
    // Test: validates required inputs and outputs
    test('validates required inputs and outputs', async () => {
      const setDefinition = {
        name: 'test-set',
        requiredInput: ['input1'],
        requiredOutput: ['task1'],
        tasks: [{ useTask: 'task1' }]
      };
      executor.setDefinition = setDefinition;

      mockRegistry.loadTask.mockResolvedValue({
        task: { name: 'task1' },
        content: 'Test content'
      });
      mockLLMClient.completeMessages.mockResolvedValue('output1_result');

      await expect(executor.execute({}, null))
        .rejects.toThrow('Set [test-set] missing required input variables: input1');

      await expect(executor.execute({ input1: 'value' }, null))
        .resolves.toBeDefined();
        
      const result = await executor.execute({ input1: 'value' }, null);
      expect(result.outputs.task1).toBe('output1_result');
    });
  });
});