ThothII/backend/test/catalog-sensitivity-analysis.test.ts

import { expect, test, vi } from "vitest";
import {
  SensitivityAnalysisInterruptedError,
  SensitivityAnalysisService,
} from "../src/catalog/sensitivity-analysis-service.js";
import { SensitivityAnalysisRunner } from "../src/catalog/sensitivity-analysis-runner.js";
import type { SensitivityClassifier } from "../src/catalog/sensitivity-classifier.js";
import type {
  CatalogColumn,
  CatalogRepository,
  CatalogTable,
  SensitivityAnalysisRun,
  WorkspaceDatabase,
} from "../src/catalog/types.js";

const database = {
  id: "11111111-1111-4111-8111-111111111111",
  workspaceId: "psd-clinical",
  engine: "postgres",
  databaseName: "warehouse",
  schema: "public",
  version: 1,
  createdAt: "2026-09-02T08:00:00Z",
  updatedAt: "2026-09-02T08:00:00Z",
  connectionStatus: "reachable",
  binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
} satisfies WorkspaceDatabase;

function catalogTable(id: string, name: string): CatalogTable {
  return {
    id,
    databaseId: database.id,
    name,
    sourceComment: null,
    description: null,
    generatedDescription: null,
    lastSyncedDatabaseVersion: 1,
    lastSyncedAt: "2026-09-02T08:00:00Z",
    version: 1,
    createdAt: "2026-09-02T08:00:00Z",
    updatedAt: "2026-09-02T08:00:00Z",
  };
}

function catalogColumn(id: string, tableId: string, name: string): CatalogColumn {
  return {
    id,
    tableId,
    name,
    ordinalPosition: 1,
    dataType: "text",
    isNullable: true,
    defaultExpression: null,
    primaryKeyPosition: null,
    isPrimaryKey: false,
    isForeignKey: false,
    foreignKeyCount: 0,
    sourceComment: null,
    description: null,
    generatedDescription: null,
    sensitive: false,
    lastSyncedDatabaseVersion: 1,
    lastSyncedAt: "2026-09-02T08:00:00Z",
    version: 1,
    createdAt: "2026-09-02T08:00:00Z",
    updatedAt: "2026-09-02T08:00:00Z",
  };
}

const running: SensitivityAnalysisRun = {
  id: "22222222-2222-4222-8222-222222222222",
  databaseId: database.id,
  scope: "all",
  engine: "local",
  modelId: null,
  policyVersion: "sensitivity-v4",
  status: "running",
  total: 0,
  suggestedSensitive: 0,
  suggestedNonSensitive: 0,
  unknown: 0,
  inputTokens: 0,
  cacheReadTokens: 0,
  outputTokens: 0,
  createdAt: "2026-09-02T08:00:00Z",
  startedAt: "2026-09-02T08:00:00Z",
  updatedAt: "2026-09-02T08:00:00Z",
  finishedAt: null,
  errorSummary: null,
};

test("stops catalog selection when the request expires during a catalog read", async () => {
  const controller = new AbortController();
  const listTables = vi.fn();
  const repository = {
    get: vi.fn(async () => {
      controller.abort();
      return database;
    }),
    listTables,
  } as unknown as CatalogRepository;
  const classifier = { assess: vi.fn() } as unknown as SensitivityClassifier;
  const analysis = new SensitivityAnalysisService(repository, classifier);

  await expect(analysis.analyze(
    database.id,
    "all",
    [],
    controller.signal,
  )).rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
  expect(listTables).not.toHaveBeenCalled();
  expect(classifier.assess).not.toHaveBeenCalled();
});

test("classifies all selected tables in one breadth-first run and reports coverage", async () => {
  const firstTable = catalogTable("33333333-3333-4333-8333-333333333333", "patients");
  const secondTable = catalogTable("44444444-4444-4444-8444-444444444444", "encounters");
  const firstColumn = catalogColumn(
    "55555555-5555-4555-8555-555555555555",
    firstTable.id,
    "status",
  );
  const secondColumn = catalogColumn(
    "66666666-6666-4666-8666-666666666666",
    secondTable.id,
    "note",
  );
  const repository = {
    get: vi.fn(async () => database),
    listTables: vi.fn(async () => [firstTable, secondTable]),
    listColumns: vi.fn(async (_databaseId: string, tableId: string) => (
      tableId === firstTable.id ? [firstColumn] : [secondColumn]
    )),
  } as unknown as CatalogRepository;
  const assess = vi.fn(async (
    _targets,
    _signal,
    _nerBudget,
    onActivity?: (message: string) => void | Promise<void>,
  ) => {
    await onActivity?.("Scanning source data: pass 1 of 3, table batch 1 of 1.");
    return [
      {
        columnId: firstColumn.id,
        assessment: "non_sensitive" as const,
        proposedSensitive: false,
        evidence: [{ kind: "coverage" as const, ruleId: "coverage.sampled_1000" }],
        observedValues: 1_000,
        coverage: "sampled" as const,
      },
      {
        columnId: secondColumn.id,
        assessment: "sensitive" as const,
        proposedSensitive: true,
        evidence: [{ kind: "content" as const, ruleId: "pii.email" }],
        observedValues: 12,
        coverage: "sampled" as const,
      },
    ];
  });
  const classifier = { assess } as unknown as SensitivityClassifier;
  const onPrepared = vi.fn();
  const onProgress = vi.fn();
  const onActivity = vi.fn();

  const suggestions = await new SensitivityAnalysisService(repository, classifier).analyze(
    database.id,
    "all",
    [],
    new AbortController().signal,
    onPrepared,
    onProgress,
    onActivity,
  );

  expect(assess).toHaveBeenCalledOnce();
  expect(assess.mock.calls[0]![0]).toEqual([
    { database, table: firstTable, columns: [firstColumn] },
    { database, table: secondTable, columns: [secondColumn] },
  ]);
  expect(onPrepared).toHaveBeenCalledWith(2);
  expect(onActivity).toHaveBeenCalledWith(
    "Scanning source data: pass 1 of 3, table batch 1 of 1.",
  );
  expect(onProgress.mock.calls.map(([processed]) => processed)).toEqual([1, 2]);
  expect(suggestions).toEqual([
    expect.objectContaining({ columnId: firstColumn.id, sensitive: false, coverage: "sampled" }),
    expect.objectContaining({ columnId: secondColumn.id, sensitive: true, coverage: "sampled" }),
  ]);
});

test("persists classifier activity in the running analysis event log", async () => {
  let persisted = running;
  const appendEvent = vi.fn(async () => undefined);
  const repository = {
    get: vi.fn(async () => database),
    createSensitivityAnalysisRun: vi.fn(async () => running),
    getSensitivityAnalysisRun: vi.fn(async () => persisted),
    updateSensitivityAnalysisRun: vi.fn(async (
      _runId: string,
      changes: Partial<SensitivityAnalysisRun>,
    ) => {
      persisted = { ...persisted, ...changes };
      return persisted;
    }),
    appendSensitivityAnalysisEvent: appendEvent,
  } as unknown as CatalogRepository;
  const analysis = {
    analyze: vi.fn(async (
      _databaseId,
      _scope,
      _targetIds,
      _signal,
      onPrepared,
      _onProgress,
      onActivity,
    ) => {
      await onPrepared?.(0);
      await onActivity?.("Scanning source data: pass 1 of 3, table batch 1 of 1.");
      return [];
    }),
  } as unknown as SensitivityAnalysisService;

  await new SensitivityAnalysisRunner(repository, analysis).run(
    database.id,
    "all",
    [],
    new AbortController().signal,
  );

  expect(appendEvent).toHaveBeenCalledWith(
    running.id,
    "info",
    "Scanning source data: pass 1 of 3, table batch 1 of 1.",
  );
});

test("marks a created run interrupted if the request deadline expires during persistence", async () => {
  const controller = new AbortController();
  const update = vi.fn(async (_runId: string, changes: Partial<SensitivityAnalysisRun>) => ({
    ...running,
    ...changes,
  }));
  const repository = {
    get: vi.fn(async () => database),
    createSensitivityAnalysisRun: vi.fn(async () => {
      controller.abort();
      return running;
    }),
    updateSensitivityAnalysisRun: update,
    appendSensitivityAnalysisEvent: vi.fn(async () => undefined),
  } as unknown as CatalogRepository;
  const analysis = { analyze: vi.fn() } as unknown as SensitivityAnalysisService;
  const runner = new SensitivityAnalysisRunner(repository, analysis);

  await expect(runner.run(database.id, "all", [], controller.signal))
    .rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
  expect(analysis.analyze).not.toHaveBeenCalled();
  expect(update).toHaveBeenCalledWith(running.id, expect.objectContaining({
    status: "interrupted",
    total: 0,
    unknown: 0,
    errorSummary: "Local sensitivity analysis was interrupted before completion.",
  }));
});