promptfoo--promptfoo/test/commands/mcp/lib/utils.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import {
  assertNotNull,
  createToolResponse,
  debounce,
  filterNonNull,
  formatDuration,
  getProperty,
  hasKey,
  retry,
  safeStringify,
  truncateText,
  withTimeout,
} from '../../../../src/commands/mcp/lib/utils';

describe('MCP Utility Functions', () => {
  describe('createToolResponse', () => {
    it('should create successful tool response with data', () => {
      const response = createToolResponse('test_tool', true, { result: 'success' });

      expect(response.isError).toBe(false);
      expect(response.content).toHaveLength(1);
      expect(response.content[0].type).toBe('text');

      const parsedContent = JSON.parse(response.content[0].text);
      expect(parsedContent.tool).toBe('test_tool');
      expect(parsedContent.success).toBe(true);
      expect(parsedContent.data).toEqual({ result: 'success' });
      expect(parsedContent.timestamp).toBeDefined();
      expect(parsedContent.error).toBeUndefined();
    });

    it('should create error tool response with error message', () => {
      const response = createToolResponse('test_tool', false, undefined, 'Something went wrong');

      expect(response.isError).toBe(true);
      expect(response.content).toHaveLength(1);

      const parsedContent = JSON.parse(response.content[0].text);
      expect(parsedContent.tool).toBe('test_tool');
      expect(parsedContent.success).toBe(false);
      expect(parsedContent.error).toBe('Something went wrong');
      expect(parsedContent.data).toBeUndefined();
    });

    it('should create response without data or error', () => {
      const response = createToolResponse('test_tool', true);

      const parsedContent = JSON.parse(response.content[0].text);
      expect(parsedContent.tool).toBe('test_tool');
      expect(parsedContent.success).toBe(true);
      expect(parsedContent.data).toBeUndefined();
      expect(parsedContent.error).toBeUndefined();
    });

    it('should include timestamp in ISO format', () => {
      const response = createToolResponse('test_tool', true);
      const parsedContent = JSON.parse(response.content[0].text);

      expect(parsedContent.timestamp).toMatch(/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$/);
    });
  });

  describe('withTimeout', () => {
    beforeEach(() => {
      vi.useFakeTimers();
    });

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

    it('should resolve when promise resolves before timeout', async () => {
      const promise = Promise.resolve('success');
      const result = await withTimeout(promise, 1000, 'Timed out');

      expect(result).toBe('success');
    });

    it('should reject when promise times out', async () => {
      const promise = new Promise((resolve) => setTimeout(() => resolve('late'), 100));
      const result = withTimeout(promise, 50, 'Operation timed out');
      const expectation = expect(result).rejects.toThrow('Operation timed out');

      await vi.runAllTimersAsync();

      await expectation;
    });

    it('should reject with original error if promise rejects', async () => {
      const promise = Promise.reject(new Error('Original error'));

      await expect(withTimeout(promise, 1000, 'Timed out')).rejects.toThrow('Original error');
    });
  });

  describe('safeStringify', () => {
    it('should stringify normal objects', () => {
      const obj = { name: 'test', value: 123 };
      const result = safeStringify(obj);

      expect(result).toBe('{"name":"test","value":123}');
    });

    it('should handle BigInt values', () => {
      const obj = { bigNumber: BigInt(123456789012345) };
      const result = safeStringify(obj);

      expect(result).toBe('{"bigNumber":"123456789012345"}');
    });

    it('should handle function values', () => {
      const obj = { func: () => 'test' };
      const result = safeStringify(obj);

      expect(result).toBe('{"func":"[Function]"}');
    });

    it('should handle undefined values', () => {
      const obj = { undef: undefined };
      const result = safeStringify(obj);

      expect(result).toBe('{"undef":"[Undefined]"}');
    });

    it('should use custom replacer when provided', () => {
      const obj = { secret: 'password', public: 'data' };
      const result = safeStringify(obj, 2, (key, value) =>
        key === 'secret' ? '[REDACTED]' : value,
      );

      expect(result).toContain('[REDACTED]');
      expect(result).toContain('data');
    });

    it('should format with spacing when provided', () => {
      const obj = { a: 1, b: 2 };
      const result = safeStringify(obj, 2);

      expect(result).toContain('\n');
      expect(result).toContain('  ');
    });
  });

  describe('hasKey', () => {
    it('should return true when key exists', () => {
      const obj = { name: 'test', value: 123 };
      const result = hasKey(obj, 'name');

      expect(result).toBe(true);
    });

    it('should return false when key does not exist', () => {
      const obj = { name: 'test' };
      const result = hasKey(obj, 'missing');

      expect(result).toBe(false);
    });

    it('should work with inherited properties', () => {
      const obj = Object.create({ inherited: 'value' });
      obj.own = 'value';

      expect(hasKey(obj, 'own')).toBe(true);
      expect(hasKey(obj, 'inherited')).toBe(true);
    });
  });

  describe('getProperty', () => {
    it('should return property value when it exists', () => {
      const obj = { name: 'test', value: 123 };
      const result = getProperty(obj, 'name', 'default');

      expect(result).toBe('test');
    });

    it('should return default value when property is undefined', () => {
      const obj = { name: 'test' };
      const result = getProperty(obj, 'missing', 'default');

      expect(result).toBe('default');
    });

    it('should return property value when it is null', () => {
      const obj = { value: null };
      const result = getProperty(obj, 'value', 'default');

      expect(result).toBeNull();
    });

    it('should return property value when it is falsy but not undefined', () => {
      const obj = { value: 0, empty: '', bool: false };

      expect(getProperty(obj, 'value', 'default')).toBe(0);
      expect(getProperty(obj, 'empty', 'default')).toBe('');
      expect(getProperty(obj, 'bool', 'default')).toBe(false);
    });
  });

  describe('assertNotNull', () => {
    it('should return value when it is not null or undefined', () => {
      expect(assertNotNull('test')).toBe('test');
      expect(assertNotNull(0)).toBe(0);
      expect(assertNotNull(false)).toBe(false);
      expect(assertNotNull('')).toBe('');
    });

    it('should throw error when value is null', () => {
      expect(() => assertNotNull(null)).toThrow('Value cannot be null or undefined');
    });

    it('should throw error when value is undefined', () => {
      expect(() => assertNotNull(undefined)).toThrow('Value cannot be null or undefined');
    });

    it('should throw custom error message', () => {
      expect(() => assertNotNull(null, 'Custom error message')).toThrow('Custom error message');
    });
  });

  describe('filterNonNull', () => {
    it('should filter out null and undefined values', () => {
      const array = ['a', null, 'b', undefined, 'c', null];
      const result = filterNonNull(array);

      expect(result).toEqual(['a', 'b', 'c']);
    });

    it('should preserve falsy but non-null values', () => {
      const array = ['', 0, false, null, undefined, 'test'];
      const result = filterNonNull(array);

      expect(result).toEqual(['', 0, false, 'test']);
    });

    it('should return empty array for all null/undefined', () => {
      const array = [null, undefined, null];
      const result = filterNonNull(array);

      expect(result).toEqual([]);
    });

    it('should return same array if no null/undefined values', () => {
      const array = ['a', 'b', 'c'];
      const result = filterNonNull(array);

      expect(result).toEqual(['a', 'b', 'c']);
    });
  });

  describe('debounce', () => {
    beforeEach(() => {
      vi.useFakeTimers();
    });

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

    it('should debounce function calls', () => {
      const mockFn = vi.fn();
      const debouncedFn = debounce(mockFn, 100);

      debouncedFn('arg1');
      debouncedFn('arg2');
      debouncedFn('arg3');

      expect(mockFn).not.toHaveBeenCalled();

      vi.advanceTimersByTime(100);

      expect(mockFn).toHaveBeenCalledTimes(1);
      expect(mockFn).toHaveBeenCalledWith('arg3');
    });

    it('should reset timer on subsequent calls', () => {
      const mockFn = vi.fn();
      const debouncedFn = debounce(mockFn, 100);

      debouncedFn('arg1');
      vi.advanceTimersByTime(50);

      debouncedFn('arg2');
      vi.advanceTimersByTime(50);

      expect(mockFn).not.toHaveBeenCalled();

      vi.advanceTimersByTime(50);

      expect(mockFn).toHaveBeenCalledTimes(1);
      expect(mockFn).toHaveBeenCalledWith('arg2');
    });
  });

  describe('retry', () => {
    it('should succeed on first attempt', async () => {
      const operation = vi.fn().mockResolvedValue('success');
      const result = await retry(operation);

      expect(result).toBe('success');
      expect(operation).toHaveBeenCalledTimes(1);
    });

    it('should retry on failure and eventually succeed', async () => {
      const operation = vi
        .fn()
        .mockRejectedValueOnce(new Error('fail1'))
        .mockRejectedValueOnce(new Error('fail2'))
        .mockResolvedValue('success');

      const result = await retry(operation, { maxAttempts: 3, baseDelay: 10 });

      expect(result).toBe('success');
      expect(operation).toHaveBeenCalledTimes(3);
    });

    it('should throw last error after max attempts', async () => {
      const operation = vi.fn().mockRejectedValue(new Error('persistent failure'));

      await expect(retry(operation, { maxAttempts: 2, baseDelay: 10 })).rejects.toThrow(
        'persistent failure',
      );

      expect(operation).toHaveBeenCalledTimes(2);
    });

    it('should use exponential backoff', async () => {
      const operation = vi
        .fn()
        .mockRejectedValueOnce(new Error('fail1'))
        .mockRejectedValueOnce(new Error('fail2'))
        .mockResolvedValue('success');

      const startTime = Date.now();

      // Mock Date.now to control timing
      const mockDateNow = vi.spyOn(Date, 'now');
      let timeOffset = 0;
      mockDateNow.mockImplementation(function () {
        return startTime + timeOffset;
      });

      // Mock setTimeout to simulate delays
      const mockSetTimeout = vi.spyOn(global, 'setTimeout');
      mockSetTimeout.mockImplementation(function (callback, delay) {
        timeOffset += delay ?? 0;
        callback();
        return {} as any;
      });

      await retry(operation, { maxAttempts: 3, baseDelay: 100, backoffFactor: 2 });

      expect(mockSetTimeout).toHaveBeenCalledWith(expect.any(Function), 100);
      expect(mockSetTimeout).toHaveBeenCalledWith(expect.any(Function), 200);

      mockDateNow.mockRestore();
      mockSetTimeout.mockRestore();
    });

    it('should respect max delay', async () => {
      const operation = vi
        .fn()
        .mockRejectedValueOnce(new Error('fail1'))
        .mockResolvedValue('success');

      const mockSetTimeout = vi.spyOn(global, 'setTimeout');
      mockSetTimeout.mockImplementation(function (callback) {
        callback();
        return {} as any;
      });

      await retry(operation, {
        maxAttempts: 2,
        baseDelay: 1000,
        backoffFactor: 10,
        maxDelay: 500,
      });

      expect(mockSetTimeout).toHaveBeenCalledWith(expect.any(Function), 500);

      mockSetTimeout.mockRestore();
    });
  });

  describe('formatDuration', () => {
    it('should format milliseconds', () => {
      expect(formatDuration(500)).toBe('500ms');
      expect(formatDuration(999)).toBe('999ms');
    });

    it('should format seconds', () => {
      expect(formatDuration(1000)).toBe('1s');
      expect(formatDuration(30000)).toBe('30s');
      expect(formatDuration(59000)).toBe('59s');
    });

    it('should format minutes and seconds', () => {
      expect(formatDuration(60000)).toBe('1m');
      expect(formatDuration(90000)).toBe('1m 30s');
      expect(formatDuration(3599000)).toBe('59m 59s');
    });

    it('should format hours and minutes', () => {
      expect(formatDuration(3600000)).toBe('1h');
      expect(formatDuration(3660000)).toBe('1h 1m');
      expect(formatDuration(7200000)).toBe('2h');
    });

    it('should handle zero duration', () => {
      expect(formatDuration(0)).toBe('0ms');
    });
  });

  describe('truncateText', () => {
    it('should return original text if shorter than max length', () => {
      const text = 'Short text';
      const result = truncateText(text, 20);

      expect(result).toBe('Short text');
    });

    it('should return original text if equal to max length', () => {
      const text = 'Exact length';
      const result = truncateText(text, 12);

      expect(result).toBe('Exact length');
    });

    it('should truncate and add ellipsis if longer than max length', () => {
      const text = 'This is a very long text that should be truncated';
      const result = truncateText(text, 20);

      expect(result).toBe('This is a very lo...');
      expect(result).toHaveLength(20);
    });

    it('should truncate without ellipsis when max length is 3 or less', () => {
      expect(truncateText('Hello', 3)).toBe('Hel');
      expect(truncateText('Hello', 2)).toBe('He');
      expect(truncateText('Hello', 1)).toBe('H');
    });

    it('should use ellipsis starting at max length 4', () => {
      const result = truncateText('Hello World', 4);
      expect(result).toBe('H...');
      expect(result).toHaveLength(4);
    });

    it('should return empty string when max length is zero', () => {
      const result = truncateText('Hello', 0);

      expect(result).toBe('');
    });

    it('should return empty string when max length is negative', () => {
      const result = truncateText('Hello', -1);

      expect(result).toBe('');
    });

    it('should handle empty string', () => {
      const result = truncateText('', 10);

      expect(result).toBe('');
    });
  });
});