Assessor_2.0/tests/regmap/regulation-retrieval-service.test.ts

import { describe, expect, it, vi } from 'vitest';
import {
  RegulationMappingUnavailableError,
  readRegulationMappingConfig,
  type RegulationMappingConfig
} from '../../server/regmap/regulation-mapping';
import { RegulationRetrievalService, groupHitsByUnit } from '../../server/regmap/regulation-retrieval-service';
import type { EmbeddingClient } from '../../server/regmap/embedding-client';
import type { QdrantCollectionInfo, QdrantSearchHit } from '../../server/regmap/qdrant-client';

const baseConfig: RegulationMappingConfig = {
  qdrantUrl: 'http://qdrant:6333',
  collection: 'regulation_catalog_v1',
  embeddingBaseUrl: 'http://embed:8080/v1',
  embeddingModel: 'test-embed',
  embeddingDimensions: 3,
  catalogVersion: 'TEST/2024',
  requestTimeoutMs: 1000,
  retrievalChunkLimit: 20,
  candidateUnitLimit: 5,
  maximumQueryCharacters: 4000
};

function fakeEmbedding(vector: number[] = [1, 0, 0]): EmbeddingClient {
  return { embed: vi.fn(async (texts: string[]) => texts.map(() => vector)) };
}

function fakeQdrant(info: QdrantCollectionInfo, hits: QdrantSearchHit[] = []) {
  return {
    collectionInfo: vi.fn(async () => info),
    search: vi.fn(async () => hits)
  };
}

const hit = (unitId: string, score: number, extra: Record<string, unknown> = {}): QdrantSearchHit => ({
  id: `${unitId}-point`,
  score,
  payload: { unit_id: unitId, control_id: unitId, control_title: `Control ${unitId}`, chunk_id: `${unitId}#1`, ...extra }
});

describe('readRegulationMappingConfig', () => {
  it('is disabled without an embedding endpoint', () => {
    expect(readRegulationMappingConfig({})).toEqual({ state: 'disabled' });
    expect(readRegulationMappingConfig({ ASSESSOR_QDRANT_URL: 'http://qdrant:6333' })).toEqual({ state: 'disabled' });
  });

  it('is misconfigured when the embedding endpoint is set but required values are missing', () => {
    const result = readRegulationMappingConfig({ ASSESSOR_EMBEDDING_BASE_URL: 'http://embed:8080/v1' });
    expect(result.state).toBe('error');
    if (result.state === 'error') expect(result.detailCode).toBe('config_incomplete');
  });

  it('parses a complete configuration and applies defaults', () => {
    const result = readRegulationMappingConfig({
      ASSESSOR_QDRANT_URL: 'http://qdrant:6333/',
      ASSESSOR_EMBEDDING_BASE_URL: 'http://embed:8080/v1/',
      ASSESSOR_EMBEDDING_MODEL: 'bge-m3',
      ASSESSOR_EMBEDDING_DIMENSIONS: '1024',
      ASSESSOR_REGMAP_CATALOG_VERSION: 'DIN EN ISO/IEC 27002:2024-01'
    });
    expect(result.state).toBe('ready');
    if (result.state === 'ready') {
      expect(result.config.qdrantUrl).toBe('http://qdrant:6333');
      expect(result.config.embeddingBaseUrl).toBe('http://embed:8080/v1');
      expect(result.config.collection).toBe('regulation_catalog_v1');
      expect(result.config.embeddingDimensions).toBe(1024);
      expect(result.config.candidateUnitLimit).toBe(5);
    }
  });

  it('rejects an invalid embedding dimension', () => {
    const result = readRegulationMappingConfig({
      ASSESSOR_QDRANT_URL: 'http://qdrant:6333',
      ASSESSOR_EMBEDDING_BASE_URL: 'http://embed:8080/v1',
      ASSESSOR_EMBEDDING_MODEL: 'bge-m3',
      ASSESSOR_EMBEDDING_DIMENSIONS: 'not-a-number',
      ASSESSOR_REGMAP_CATALOG_VERSION: 'X'
    });
    expect(result.state).toBe('error');
  });
});

describe('RegulationRetrievalService capability', () => {
  it('reports collection_missing when the collection does not exist', async () => {
    const service = new RegulationRetrievalService(baseConfig, {
      embeddingClient: fakeEmbedding(),
      qdrantClient: fakeQdrant({ exists: false })
    });
    const report = await service.refreshCapability();
    expect(report).toMatchObject({ capability: 'misconfigured', detailCode: 'collection_missing' });
  });

  it('reports dimension_mismatch when the stored vector size differs', async () => {
    const service = new RegulationRetrievalService(baseConfig, {
      embeddingClient: fakeEmbedding(),
      qdrantClient: fakeQdrant({ exists: true, vectorSize: 768, distance: 'Cosine' })
    });
    expect(await service.refreshCapability()).toMatchObject({ capability: 'misconfigured', detailCode: 'dimension_mismatch' });
  });

  it('reports distance_mismatch when the collection is not cosine', async () => {
    const service = new RegulationRetrievalService(baseConfig, {
      embeddingClient: fakeEmbedding(),
      qdrantClient: fakeQdrant({ exists: true, vectorSize: 3, distance: 'Dot' })
    });
    expect(await service.refreshCapability()).toMatchObject({ capability: 'misconfigured', detailCode: 'distance_mismatch' });
  });

  it('is ready when the collection matches and the embedding probe succeeds', async () => {
    const service = new RegulationRetrievalService(baseConfig, {
      embeddingClient: fakeEmbedding(),
      qdrantClient: fakeQdrant({ exists: true, vectorSize: 3, distance: 'Cosine' })
    });
    expect(await service.refreshCapability()).toMatchObject({ capability: 'ready', catalogVersion: 'TEST/2024', embeddingModel: 'test-embed' });
  });

  it('surfaces an embedding failure as misconfigured', async () => {
    const embedding: EmbeddingClient = { embed: vi.fn(async () => { throw Object.assign(new Error('down'), { code: 'embedding_unreachable' }); }) };
    const service = new RegulationRetrievalService(baseConfig, {
      embeddingClient: embedding,
      qdrantClient: fakeQdrant({ exists: true, vectorSize: 3, distance: 'Cosine' })
    });
    expect(await service.refreshCapability()).toMatchObject({ capability: 'misconfigured', detailCode: 'embedding_unreachable' });
  });
});

describe('groupHitsByUnit', () => {
  it('groups chunks to their unit, scoring each unit by its best chunk', () => {
    const grouped = groupHitsByUnit([hit('A', 0.5), hit('B', 0.9), hit('A', 0.7)], 5);
    expect(grouped.map(c => [c.unitId, c.score])).toEqual([['B', 0.9], ['A', 0.7]]);
  });

  it('is deterministic on ties and honours the unit limit', () => {
    const grouped = groupHitsByUnit([hit('B', 0.8), hit('A', 0.8), hit('C', 0.4)], 2);
    expect(grouped.map(c => c.unitId)).toEqual(['A', 'B']);
  });

  it('ignores hits without a unit id', () => {
    const grouped = groupHitsByUnit([{ id: 'x', score: 0.9, payload: {} }, hit('A', 0.3)], 5);
    expect(grouped.map(c => c.unitId)).toEqual(['A']);
  });
});

describe('RegulationRetrievalService retrieve', () => {
  it('embeds the query and returns grouped candidate units', async () => {
    const embedding = fakeEmbedding([0, 1, 0]);
    const qdrant = fakeQdrant({ exists: true, vectorSize: 3, distance: 'Cosine' }, [hit('8.24', 0.91), hit('5.1', 0.42)]);
    const service = new RegulationRetrievalService(baseConfig, { embeddingClient: embedding, qdrantClient: qdrant });
    await service.refreshCapability();

    const result = await service.retrieve({ title: 'RC4 enabled on DC', description: 'Weak cipher permitted' });
    expect(result.catalogVersion).toBe('TEST/2024');
    expect(result.candidates.map(c => c.unitId)).toEqual(['8.24', '5.1']);
    expect(result.candidates[0]).toMatchObject({ controlId: '8.24', controlTitle: 'Control 8.24', bestChunkId: '8.24#1' });
    expect(embedding.embed).toHaveBeenCalledWith(expect.arrayContaining([expect.stringContaining('RC4 enabled on DC')]), 'query');
  });

  it('throws RegulationMappingUnavailableError when not ready', async () => {
    const service = new RegulationRetrievalService(baseConfig, {
      embeddingClient: fakeEmbedding(),
      qdrantClient: fakeQdrant({ exists: false })
    });
    await expect(service.retrieve({ title: 'anything' })).rejects.toBeInstanceOf(RegulationMappingUnavailableError);
  });

  it('does not log the query text', async () => {
    const logs: unknown[] = [];
    const spies = [
      vi.spyOn(console, 'log').mockImplementation((...args) => { logs.push(...args); }),
      vi.spyOn(console, 'error').mockImplementation((...args) => { logs.push(...args); }),
      vi.spyOn(console, 'warn').mockImplementation((...args) => { logs.push(...args); })
    ];
    const service = new RegulationRetrievalService(baseConfig, {
      embeddingClient: fakeEmbedding([0, 1, 0]),
      qdrantClient: fakeQdrant({ exists: true, vectorSize: 3, distance: 'Cosine' }, [hit('8.24', 0.91)])
    });
    await service.refreshCapability();
    await service.retrieve({ title: 'SECRET-MARKER-TITLE', description: 'SECRET-MARKER-DESC' });
    spies.forEach(spy => spy.mockRestore());
    expect(logs.join(' ')).not.toContain('SECRET-MARKER');
  });
});