promptfoo--promptfoo/test/redteam/plugins/policy/utils.test.ts

import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest';
import { POLICY_METRIC_PREFIX } from '../../../../src/redteam/plugins/policy/constants';
import {
  deserializePolicyIdFromMetric,
  determinePolicyTypeFromId,
  formatPolicyIdentifierAsMetric,
  isPolicyMetric,
  isValidPolicyObject,
  makeCustomPolicyCloudUrl,
  makeInlinePolicyId,
  makeInlinePolicyIdSync,
} from '../../../../src/redteam/plugins/policy/utils';
import { PolicyObjectSchema } from '../../../../src/redteam/types';
import { sha256 } from '../../../../src/util/createHash';
import { isUuid } from '../../../../src/util/uuid';

// Mock dependencies
vi.mock('../../../../src/util/uuid', () => ({
  isUuid: vi.fn(),
}));

vi.mock('../../../../src/util/createHash', () => ({
  sha256: vi.fn(),
}));

describe('Policy Utils', () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  describe('isPolicyMetric', () => {
    it('should return true for metrics that start with POLICY_METRIC_PREFIX', () => {
      expect(isPolicyMetric(`${POLICY_METRIC_PREFIX}:test`)).toBe(true);
      expect(isPolicyMetric(`${POLICY_METRIC_PREFIX}`)).toBe(true);
      expect(isPolicyMetric(`${POLICY_METRIC_PREFIX}:{"id":"123"}`)).toBe(true);
    });

    it('should return false for metrics that do not start with POLICY_METRIC_PREFIX', () => {
      expect(isPolicyMetric('test:policy')).toBe(false);
      expect(isPolicyMetric('other')).toBe(false);
      expect(isPolicyMetric('')).toBe(false);
      expect(isPolicyMetric('Policy')).toBe(false);
    });
  });

  describe('deserializePolicyIdFromMetric', () => {
    it('should deserialize a policy ID from a metric without strategy suffix', () => {
      const result = deserializePolicyIdFromMetric(`${POLICY_METRIC_PREFIX}:test-id`);
      expect(result).toBe('test-id');
    });

    it('should deserialize a policy ID and remove strategy suffix', () => {
      const mockUuid = '550e8400-e29b-41d4-a716-446655440000';

      const result = deserializePolicyIdFromMetric(`${POLICY_METRIC_PREFIX}:${mockUuid}/jailbreak`);
      expect(result).toBe(mockUuid);
    });

    it('should handle multiple strategy suffixes', () => {
      const mockUuid = '550e8400-e29b-41d4-a716-446655440000';

      const result = deserializePolicyIdFromMetric(
        `${POLICY_METRIC_PREFIX}:${mockUuid}/jailbreak/extra`,
      );
      expect(result).toBe(mockUuid);
    });

    it('should handle various strategy names', () => {
      expect(deserializePolicyIdFromMetric(`${POLICY_METRIC_PREFIX}:policy-id/GOAT`)).toBe(
        'policy-id',
      );
      expect(
        deserializePolicyIdFromMetric(`${POLICY_METRIC_PREFIX}:policy-id/prompt-injection`),
      ).toBe('policy-id');
      expect(
        deserializePolicyIdFromMetric(`${POLICY_METRIC_PREFIX}:abcdef123456/harmful-content`),
      ).toBe('abcdef123456');
    });

    it('should handle inline policy IDs with strategy suffix', () => {
      const inlineId = 'abcdef123456';
      const result = deserializePolicyIdFromMetric(`${POLICY_METRIC_PREFIX}:${inlineId}/jailbreak`);
      expect(result).toBe(inlineId);
    });
  });

  describe('formatPolicyIdentifierAsMetric', () => {
    it('should format a policy identifier correctly without originalMetric', () => {
      expect(formatPolicyIdentifierAsMetric('test-policy')).toBe('Policy: test-policy');
      expect(formatPolicyIdentifierAsMetric('123')).toBe('Policy: 123');
      expect(formatPolicyIdentifierAsMetric('')).toBe('Policy: ');
    });

    it('should handle identifiers with special characters', () => {
      expect(formatPolicyIdentifierAsMetric('policy-with-dashes')).toBe(
        'Policy: policy-with-dashes',
      );
      expect(formatPolicyIdentifierAsMetric('policy_with_underscores')).toBe(
        'Policy: policy_with_underscores',
      );
      expect(formatPolicyIdentifierAsMetric('policy.with.dots')).toBe('Policy: policy.with.dots');
    });

    it('should preserve strategy suffix from originalMetric', () => {
      const mockUuid = '550e8400-e29b-41d4-a716-446655440000';
      expect(
        formatPolicyIdentifierAsMetric('test-policy', `PolicyViolation:${mockUuid}/jailbreak`),
      ).toBe('Policy/jailbreak: test-policy');
      expect(formatPolicyIdentifierAsMetric('my-policy', `PolicyViolation:${mockUuid}/GOAT`)).toBe(
        'Policy/GOAT: my-policy',
      );
    });

    it('should handle empty originalMetric', () => {
      expect(formatPolicyIdentifierAsMetric('test-policy', '')).toBe('Policy: test-policy');
    });
  });

  describe('makeCustomPolicyCloudUrl', () => {
    it('should create a valid cloud URL', () => {
      const result = makeCustomPolicyCloudUrl('https://cloud.example.com', 'policy-123');
      expect(result).toBe('https://cloud.example.com/redteam/plugins/policies/policy-123');
    });

    it('should handle URLs with trailing slash', () => {
      const result = makeCustomPolicyCloudUrl('https://cloud.example.com/', 'policy-456');
      expect(result).toBe('https://cloud.example.com//redteam/plugins/policies/policy-456');
    });

    it('should handle various policy ID formats', () => {
      expect(makeCustomPolicyCloudUrl('https://app.com', 'uuid-123-456')).toBe(
        'https://app.com/redteam/plugins/policies/uuid-123-456',
      );
      expect(makeCustomPolicyCloudUrl('https://app.com', 'hash12345678')).toBe(
        'https://app.com/redteam/plugins/policies/hash12345678',
      );
    });

    it('should handle empty strings', () => {
      expect(makeCustomPolicyCloudUrl('', 'policy-id')).toBe('/redteam/plugins/policies/policy-id');
      expect(makeCustomPolicyCloudUrl('https://app.com', '')).toBe(
        'https://app.com/redteam/plugins/policies/',
      );
    });
  });

  describe('determinePolicyTypeFromId', () => {
    it('should return "reusable" for valid UUIDs', () => {
      (isUuid as Mock).mockReturnValue(true);
      expect(determinePolicyTypeFromId('550e8400-e29b-41d4-a716-446655440000')).toBe('reusable');
      expect(isUuid).toHaveBeenCalledWith('550e8400-e29b-41d4-a716-446655440000');
    });

    it('should return "inline" for non-UUID strings', () => {
      (isUuid as Mock).mockReturnValue(false);
      expect(determinePolicyTypeFromId('hash123456')).toBe('inline');
      expect(determinePolicyTypeFromId('not-a-uuid')).toBe('inline');
      expect(determinePolicyTypeFromId('')).toBe('inline');
    });

    it('should handle edge cases', () => {
      (isUuid as Mock).mockReturnValue(false);
      expect(determinePolicyTypeFromId('123')).toBe('inline');
      expect(determinePolicyTypeFromId('550e8400')).toBe('inline');
    });
  });

  describe('isValidPolicyObject', () => {
    // Mock the PolicyObjectSchema.safeParse method
    const mockSafeParse = vi.spyOn(PolicyObjectSchema, 'safeParse');

    it('should return true for valid PolicyObject', () => {
      mockSafeParse.mockReturnValueOnce({ success: true, data: {} as any });
      const policy = {
        id: 'test-id',
        name: 'Test Policy',
        text: 'Policy text',
      };

      expect(isValidPolicyObject(policy)).toBe(true);
      expect(mockSafeParse).toHaveBeenCalledWith(policy);
    });

    it('should return false for invalid PolicyObject', () => {
      mockSafeParse.mockReturnValueOnce({ success: false, error: {} as any });
      const invalidPolicy = {
        invalid: 'field',
      };

      expect(isValidPolicyObject(invalidPolicy as any)).toBe(false);
      expect(mockSafeParse).toHaveBeenCalledWith(invalidPolicy);
    });

    it('should return false for string policy', () => {
      mockSafeParse.mockReturnValueOnce({ success: false, error: {} as any });
      const stringPolicy = 'This is a string policy';

      expect(isValidPolicyObject(stringPolicy)).toBe(false);
      expect(mockSafeParse).toHaveBeenCalledWith(stringPolicy);
    });

    it('should handle null and undefined', () => {
      mockSafeParse.mockReturnValueOnce({ success: false, error: {} as any });
      expect(isValidPolicyObject(null as any)).toBe(false);

      mockSafeParse.mockReturnValueOnce({ success: false, error: {} as any });
      expect(isValidPolicyObject(undefined as any)).toBe(false);
    });
  });

  describe('makeInlinePolicyId', () => {
    it('should create a 12-character ID from policy text', async () => {
      (sha256 as Mock).mockResolvedValue('abcdef1234567890abcdef1234567890');
      const result = await makeInlinePolicyId('This is my policy text');

      expect(result).toBe('abcdef123456');
      expect(result).toHaveLength(12);
      expect(sha256).toHaveBeenCalledWith('This is my policy text');
    });

    it('should create consistent IDs for the same text', async () => {
      (sha256 as Mock).mockResolvedValue('1234567890abcdef1234567890abcdef');
      const text = 'Same policy text';
      const id1 = await makeInlinePolicyId(text);
      const id2 = await makeInlinePolicyId(text);

      expect(id1).toBe(id2);
      expect(id1).toBe('1234567890ab');
    });

    it('should create different IDs for different text', async () => {
      (sha256 as Mock)
        .mockResolvedValueOnce('aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa')
        .mockResolvedValueOnce('bbbbbbbbbbbbbbbbbbbbbbbbbbbbbb');

      const id1 = await makeInlinePolicyId('Policy A');
      const id2 = await makeInlinePolicyId('Policy B');

      expect(id1).toBe('aaaaaaaaaaaa');
      expect(id2).toBe('bbbbbbbbbbbb');
      expect(id1).not.toBe(id2);
    });

    it('should handle empty strings', async () => {
      (sha256 as Mock).mockResolvedValue('emptyhashabcd1234567890abcdef');
      const result = await makeInlinePolicyId('');

      expect(result).toBe('emptyhashabc');
      expect(sha256).toHaveBeenCalledWith('');
    });

    it('should handle special characters and multiline text', async () => {
      (sha256 as Mock).mockResolvedValue('specialhash1234567890abcdef');
      const text = `Line 1
Line 2
Special chars: !@#$%^&*()
Unicode: 你好 🎉`;

      const result = await makeInlinePolicyId(text);
      expect(result).toBe('specialhash1');
      expect(sha256).toHaveBeenCalledWith(text);
    });

    it('should always return 12 characters even if hash is shorter', async () => {
      (sha256 as Mock).mockResolvedValue('short');
      const result = await makeInlinePolicyId('Short hash text');

      expect(result).toBe('short');
      expect(result.length).toBeLessThanOrEqual(12);
    });
  });

  describe('makeInlinePolicyIdSync', () => {
    it('should create a 12-character ID from policy text synchronously', () => {
      (sha256 as Mock).mockReturnValue('abcdef1234567890abcdef1234567890');
      const result = makeInlinePolicyIdSync('This is my policy text');

      expect(result).toBe('abcdef123456');
      expect(result).toHaveLength(12);
      expect(sha256).toHaveBeenCalledWith('This is my policy text');
    });

    it('should create consistent IDs for the same text', () => {
      (sha256 as Mock).mockReturnValue('1234567890abcdef1234567890abcdef');
      const text = 'Same policy text';
      const id1 = makeInlinePolicyIdSync(text);
      const id2 = makeInlinePolicyIdSync(text);

      expect(id1).toBe(id2);
      expect(id1).toBe('1234567890ab');
    });

    it('should handle empty strings', () => {
      (sha256 as Mock).mockReturnValue('emptyhashabcd1234567890abcdef');
      const result = makeInlinePolicyIdSync('');

      expect(result).toBe('emptyhashabc');
      expect(sha256).toHaveBeenCalledWith('');
    });
  });
});