aurak/server/src/assessment/graph/nodes/grader.node.spec.ts

import { graderNode } from './grader.node';
import { HumanMessage, AIMessage } from '@langchain/core/messages';

function mockModel(response: any) {
  return {
    invoke: jest.fn().mockResolvedValue({
      content: JSON.stringify(response),
    }),
  };
}

function baseState(overrides: any = {}) {
  return {
    messages: [new HumanMessage('test answer')],
    questions: [{ id: 'q1', questionText: 'What is JS?', keyPoints: ['point1'], dimension: 'llm' }],
    currentQuestionIndex: 0,
    scores: {},
    feedbackHistory: [],
    followUpCount: 0,
    shouldFollowUp: false,
    questionCount: 5,
    language: 'en',
    ...overrides,
  } as any;
}

describe('graderNode', () => {
  describe('validation guards', () => {
    it('should throw when model is missing', async () => {
      await expect(graderNode(baseState(), { configurable: {} } as any)).rejects.toThrow('Missing model');
    });

    it('should return empty object when last message is not HumanMessage', async () => {
      const state = baseState({ messages: [new AIMessage('I am AI')] });
      const result = await graderNode(state, { configurable: { model: mockModel({}) } } as any);
      expect(result).toEqual({});
    });

    it('should skip question and advance index when current question not found', async () => {
      const state = baseState({ currentQuestionIndex: 99, questions: [{ id: 'q1', questionText: 'Q', keyPoints: ['k'], dimension: 'llm' }] });
      const result = await graderNode(state, { configurable: { model: mockModel({}) } } as any);
      expect(result.currentQuestionIndex).toBe(100);
    });
  });

  describe('breakout logic (shouldFollowUp overrides)', () => {
    it('should NOT follow up when followUpCount >= 2 even if LLM says follow up', async () => {
      const model = mockModel({ score: 5, feedback: 'needs work', should_follow_up: true, follow_up_question: 'More?' });
      const state = baseState({ followUpCount: 2 });
      const result = await graderNode(state, { configurable: { model } } as any);
      expect(result.shouldFollowUp).toBe(false);
    });

    it('should NOT follow up when score >= 8 even if LLM says follow up', async () => {
      const model = mockModel({ score: 9, feedback: 'good', should_follow_up: true });
      const state = baseState();
      const result = await graderNode(state, { configurable: { model } } as any);
      expect(result.shouldFollowUp).toBe(false);
    });

    it('should NOT follow up when user says "I don\'t know"', async () => {
      const model = mockModel({ score: 2, feedback: 'no answer', should_follow_up: true });
      const state = baseState({ messages: [new HumanMessage("no idea")] });
      const result = await graderNode(state, { configurable: { model } } as any);
      expect(result.shouldFollowUp).toBe(false);
    });

    it('should allow follow up when conditions are met', async () => {
      const model = mockModel({ score: 5, feedback: 'incomplete', should_follow_up: true, follow_up_question: 'Can you elaborate?' });
      const state = baseState({ followUpCount: 0 });
      const result = await graderNode(state, { configurable: { model } } as any);
      expect(result.shouldFollowUp).toBe(true);
      expect(result.followUpCount).toBe(1);
    });
  });

  describe('error handling', () => {
    it('should handle LLM returning invalid JSON gracefully', async () => {
      const model = { invoke: jest.fn().mockResolvedValue({ content: 'NOT JSON' }) };
      const result = await graderNode(baseState(), { configurable: { model } } as any);
      expect(result.currentQuestionIndex).toBe(1);
      expect(result.shouldFollowUp).toBe(false);
    });
  });

  describe('scoring and indexing', () => {
    it('should advance currentQuestionIndex when not following up', async () => {
      const model = mockModel({ score: 6, feedback: 'ok', should_follow_up: false });
      const result = await graderNode(baseState(), { configurable: { model } } as any);
      expect(result.currentQuestionIndex).toBe(1);
      expect(result.scores).toBeDefined();
    });

    it('should keep currentQuestionIndex when following up', async () => {
      const model = mockModel({ score: 5, feedback: 'needs work', should_follow_up: true, follow_up_question: 'Can you clarify?' });
      const state = baseState({ followUpCount: 0 });
      const result = await graderNode(state, { configurable: { model } } as any);
      expect(result.currentQuestionIndex).toBe(0);
    });

    it('should record score under question id in scores map', async () => {
      const model = mockModel({ score: 7, feedback: 'good', should_follow_up: false });
      const state = baseState({ questions: [{ id: 'q-test', questionText: 'Q', keyPoints: ['k'], dimension: 'llm' }] });
      const result = await graderNode(state, { configurable: { model } } as any);
      expect((result.scores as any)['q-test']).toBe(7);
    });
  });

  describe('language support', () => {
    it('should handle Chinese language', async () => {
      const model = mockModel({ score: 8, feedback: '很好', should_follow_up: false });
      const state = baseState({ language: 'zh' });
      const result = await graderNode(state, { configurable: { model } } as any);
      expect(result).toBeDefined();
    });

    it('should handle Japanese language', async () => {
      const model = mockModel({ score: 8, feedback: '良い', should_follow_up: false });
      const state = baseState({ language: 'ja' });
      const result = await graderNode(state, { configurable: { model } } as any);
      expect(result).toBeDefined();
    });
  });
});