promptfoo--promptfoo/test/assertions/contextUtils.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { resolveContext } from '../../src/assertions/contextUtils';
import * as transformUtil from '../../src/util/transform';

import type { Assertion, AtomicTestCase } from '../../src/types/index';

vi.mock('../../src/util/transform', async (importOriginal) => {
  const actual = await importOriginal<typeof transformUtil>();
  return {
    ...actual,
    transform: vi.fn(),
  };
});

describe('resolveContext', () => {
  const mockTransform = vi.mocked(transformUtil.transform);

  beforeEach(() => {
    vi.clearAllMocks();
    mockTransform.mockReset();
  });

  afterEach(() => {
    mockTransform.mockReset();
    vi.clearAllMocks();
  });

  describe('with context variable', () => {
    it('should return context from test.vars.context', async () => {
      const assertion: Assertion = { type: 'context-faithfulness' };
      const test: AtomicTestCase = {
        vars: { context: 'test context' },
        options: {},
      };

      const result = await resolveContext(assertion, test, 'output', 'prompt');
      expect(result).toBe('test context');
      expect(mockTransform).not.toHaveBeenCalled();
    });

    it('should return context array from test.vars.context', async () => {
      const assertion: Assertion = { type: 'context-faithfulness' };
      const contextArray = ['chunk 1', 'chunk 2', 'chunk 3'];
      const test: AtomicTestCase = {
        vars: { context: contextArray },
        options: {},
      };

      const result = await resolveContext(assertion, test, 'output', 'prompt');
      expect(result).toEqual(contextArray);
      expect(mockTransform).not.toHaveBeenCalled();
    });

    it('should reject context array with non-string elements', async () => {
      const assertion: Assertion = { type: 'context-faithfulness' };
      const test: AtomicTestCase = {
        vars: { context: ['valid', 123, 'invalid'] },
        options: {},
      };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        'Invalid context: expected an array of strings, but found number at index 1',
      );
    });

    it('should use fallback context when no context variable', async () => {
      const assertion: Assertion = { type: 'context-recall' };
      const test: AtomicTestCase = { vars: {}, options: {} };

      const result = await resolveContext(assertion, test, 'output', 'prompt', 'fallback context');
      expect(result).toBe('fallback context');
      expect(mockTransform).not.toHaveBeenCalled();
    });
  });

  describe('with contextTransform', () => {
    it('should transform output to get context', async () => {
      mockTransform.mockResolvedValue('transformed context' as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.context',
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      const result = await resolveContext(assertion, test, { context: 'data' }, 'prompt');

      expect(result).toBe('transformed context');
      expect(mockTransform).toHaveBeenCalledWith(
        'output.context',
        { context: 'data' },
        {
          vars: {},
          prompt: { label: 'prompt' },
        },
      );
    });

    it('should transform output to get context array', async () => {
      const mockArray = ['chunk 1', 'chunk 2', 'chunk 3'];
      mockTransform.mockResolvedValue(mockArray as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.chunks',
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      const result = await resolveContext(assertion, test, { chunks: ['a', 'b', 'c'] }, 'prompt');

      expect(result).toEqual(mockArray);
      expect(mockTransform).toHaveBeenCalledWith(
        'output.chunks',
        { chunks: ['a', 'b', 'c'] },
        {
          vars: {},
          prompt: { label: 'prompt' },
        },
      );
    });

    it('should prioritize contextTransform over context variable', async () => {
      mockTransform.mockResolvedValue('transformed context' as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.context',
      };
      const test: AtomicTestCase = {
        vars: { context: 'original context' },
        options: {},
      };

      const result = await resolveContext(assertion, test, { context: 'data' }, 'prompt');

      expect(result).toBe('transformed context');
      expect(mockTransform).toHaveBeenCalledWith(
        'output.context',
        { context: 'data' },
        {
          vars: { context: 'original context' },
          prompt: { label: 'prompt' },
        },
      );
    });

    it('should throw error if transform returns non-string and non-string-array', async () => {
      mockTransform.mockResolvedValue(123 as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.invalid',
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        `contextTransform must return a string or array of strings. Got number. Check your transform expression: ${transformUtil.INLINE_STRING_LABEL}: output.invalid`,
      );
    });

    it('should throw error if transform returns array with non-string elements', async () => {
      mockTransform.mockResolvedValue(['valid', 123, 'invalid'] as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.invalid',
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        `contextTransform must return a string or array of strings. Got object. Check your transform expression: ${transformUtil.INLINE_STRING_LABEL}: output.invalid`,
      );
    });

    it('should throw error if transform fails', async () => {
      mockTransform.mockRejectedValue(new Error('Transform failed'));

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.invalid',
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        `Failed to transform context using expression '${transformUtil.INLINE_STRING_LABEL}: output.invalid': Transform failed`,
      );
    });

    it('should support inline function as contextTransform', async () => {
      const contextFn = (output: any) => output.context;
      mockTransform.mockResolvedValue('extracted context' as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: contextFn,
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      const result = await resolveContext(
        assertion,
        test,
        { context: 'extracted context', data: 'other' },
        'prompt',
      );

      expect(result).toBe('extracted context');
      expect(mockTransform).toHaveBeenCalledWith(
        contextFn,
        { context: 'extracted context', data: 'other' },
        expect.any(Object),
      );
    });

    it('should show [inline function] in error messages for function contextTransform', async () => {
      const contextFn = (() => 42) as any;
      mockTransform.mockRejectedValue(new Error('Transform failed'));

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: contextFn,
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        /Failed to transform context using expression '\[inline function\].*': Transform failed/,
      );
    });
  });

  describe('error cases', () => {
    it('should throw error when no context is available', async () => {
      const assertion: Assertion = { type: 'context-faithfulness' };
      const test: AtomicTestCase = { vars: {}, options: {} };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        'Context is required for context-based assertions. Provide either a "context" variable (string or array of strings) in your test case or use "contextTransform" to extract context from the provider response.',
      );
    });

    it('should throw error when context is empty string', async () => {
      const assertion: Assertion = { type: 'context-faithfulness' };
      const test: AtomicTestCase = {
        vars: { context: '' },
        options: {},
      };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        'Context is required for context-based assertions. Provide either a "context" variable (string or array of strings) in your test case or use "contextTransform" to extract context from the provider response.',
      );
    });

    it('should throw error when context is empty array', async () => {
      const assertion: Assertion = { type: 'context-faithfulness' };
      const test: AtomicTestCase = {
        vars: { context: [] },
        options: {},
      };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        'Context is required for context-based assertions. Provide either a "context" variable (string or array of strings) in your test case or use "contextTransform" to extract context from the provider response.',
      );
    });

    it('should throw error when context array contains empty strings', async () => {
      const assertion: Assertion = { type: 'context-faithfulness' };
      const test: AtomicTestCase = {
        vars: { context: ['valid chunk', '', 'another chunk'] },
        options: {},
      };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        'Context is required for context-based assertions. Provide either a "context" variable (string or array of strings) in your test case or use "contextTransform" to extract context from the provider response.',
      );
    });

    it('should throw error when contextTransform returns empty string', async () => {
      mockTransform.mockResolvedValue('' as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.empty',
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        'Context is required for context-based assertions. Provide either a "context" variable (string or array of strings) in your test case or use "contextTransform" to extract context from the provider response.',
      );
    });

    it('should throw error when contextTransform returns empty array', async () => {
      mockTransform.mockResolvedValue([] as any);

      const assertion: Assertion = {
        type: 'context-faithfulness',
        contextTransform: 'output.empty',
      };
      const test: AtomicTestCase = { vars: {}, options: {} };

      await expect(resolveContext(assertion, test, 'output', 'prompt')).rejects.toThrow(
        'Context is required for context-based assertions. Provide either a "context" variable (string or array of strings) in your test case or use "contextTransform" to extract context from the provider response.',
      );
    });
  });
});