wulf-pulse/lib/services/remediation-service.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';

// Mock postgresClient and phishing-audit BEFORE importing the module under test.
const queryMock = vi.fn();
const transactionMock = vi.fn();
vi.mock('./postgres-client', () => ({
  postgresClient: {
    query: (...args: unknown[]) => queryMock(...args),
    transaction: (...args: unknown[]) => transactionMock(...args),
  },
}));

const writeAuditEventMock = vi.fn();
vi.mock('./phishing-audit', () => ({
  writeAuditEvent: (...args: unknown[]) => writeAuditEventMock(...args),
}));

const generateAndPostAcknowledgmentMock = vi.fn();
vi.mock('./triage-note-service', () => ({
  generateAndPostAcknowledgment: (...args: unknown[]) => generateAndPostAcknowledgmentMock(...args),
}));

// eslint-disable-next-line import/first -- imported after vi.mock hoisting
import {
  approveRemediationActions,
  remediateApprovedActions,
  markCampaignFalsePositive,
  RemediationValidationError,
  RemediationConflictError,
} from './remediation-service';

/** Records every SQL call made through the fake transaction client. */
let clientCalls: Array<{ sql: string; params: unknown[] }> = [];

interface MockRows {
  classification?: unknown[];
  insertRemediation?: unknown[];
  remediationRows?: unknown[];
  campaign?: unknown[];
  guardRows?: unknown[];
}

function makeClient(rows: MockRows) {
  return {
    query: vi.fn(async (sql: string, params?: unknown[]) => {
      clientCalls.push({ sql, params: params ?? [] });

      if (sql.includes('FROM classifications')) {
        return { rows: rows.classification ?? [], rowCount: rows.classification?.length ?? 0 };
      }
      if (sql.includes('INSERT INTO remediation_actions')) {
        const insertRows = rows.insertRemediation ?? [];
        return { rows: [insertRows.shift() ?? { id: 'unstaged-action-id' }], rowCount: 1 };
      }
      if (sql.includes('SELECT id::text, action_type, status') && sql.includes('FOR UPDATE')) {
        return { rows: rows.remediationRows ?? [], rowCount: rows.remediationRows?.length ?? 0 };
      }
      if (sql.includes('UPDATE remediation_actions SET status')) {
        return { rows: [], rowCount: 1 };
      }
      if (sql.includes('FROM remediation_actions') && sql.includes('status IN')) {
        return { rows: rows.guardRows ?? [], rowCount: rows.guardRows?.length ?? 0 };
      }
      if (sql.includes('SELECT status FROM campaigns')) {
        return { rows: rows.campaign ?? [], rowCount: rows.campaign?.length ?? 0 };
      }
      if (sql.includes('UPDATE campaigns SET status')) {
        return { rows: [], rowCount: 1 };
      }
      throw new Error(`Unstaged query in test mock: ${sql}`);
    }),
  };
}

function callsContaining(needle: string) {
  return clientCalls.filter((c) => c.sql.includes(needle));
}

function stage(rows: MockRows) {
  transactionMock.mockImplementation(async (callback: (client: unknown) => Promise<unknown>) =>
    callback(makeClient(rows))
  );
}

beforeEach(() => {
  queryMock.mockReset();
  transactionMock.mockReset();
  writeAuditEventMock.mockReset();
  writeAuditEventMock.mockResolvedValue('audit-id');
  generateAndPostAcknowledgmentMock.mockReset();
  generateAndPostAcknowledgmentMock.mockResolvedValue({ noteText: 'thanks', tickets: [] });
  clientCalls = [];
});

// =============================================================================
// approveRemediationActions
// =============================================================================

describe('approveRemediationActions', () => {
  it('rejects an action type that is NOT in the latest classification recommended_actions', async () => {
    stage({
      classification: [{ recommended_actions: ['warn_user'] }],
    });

    await expect(
      approveRemediationActions('campaign-1', [{ actionType: 'block_sender' }], 'operator@example.com')
    ).rejects.toThrow(RemediationValidationError);
  });

  it('throws RemediationValidationError when no classification row exists for the campaign', async () => {
    stage({ classification: [] });

    await expect(
      approveRemediationActions('campaign-1', [{ actionType: 'warn_user' }], 'operator@example.com')
    ).rejects.toThrow(RemediationValidationError);
  });

  it('inserts one remediation_actions row per requested action plus one audit row, inside one transaction', async () => {
    stage({
      classification: [{ recommended_actions: ['block_sender', 'purge_message'] }],
      insertRemediation: [{ id: 'action-1' }, { id: 'action-2' }],
    });

    const result = await approveRemediationActions(
      'campaign-1',
      [
        { actionType: 'block_sender', params: { sender: 'evil@example.com' } },
        { actionType: 'purge_message', params: {} },
      ],
      'operator@example.com'
    );

    expect(transactionMock).toHaveBeenCalledTimes(1);
    const insertCalls = callsContaining('INSERT INTO remediation_actions');
    expect(insertCalls).toHaveLength(2);
    expect(insertCalls[0].sql).toContain("'approved'");
    expect(insertCalls[0].params).toEqual([
      'campaign-1',
      'block_sender',
      JSON.stringify({ sender: 'evil@example.com' }),
      'operator@example.com',
    ]);

    expect(writeAuditEventMock).toHaveBeenCalledTimes(1);
    const [auditArgs, auditClient] = writeAuditEventMock.mock.calls[0];
    expect(auditArgs).toMatchObject({
      campaignId: 'campaign-1',
      actor: 'operator@example.com',
      eventType: 'remediation_approved',
    });
    expect(auditClient).toBeDefined();

    expect(result).toHaveLength(2);
    expect(result[0]).toMatchObject({ id: 'action-1', actionType: 'block_sender', status: 'approved' });
  });
});

// =============================================================================
// remediateApprovedActions
// =============================================================================

describe('remediateApprovedActions', () => {
  it('transitions every status=approved row to completed and writes one audit row per transitioned action', async () => {
    stage({
      remediationRows: [
        { id: 'action-1', action_type: 'block_sender', status: 'approved' },
        { id: 'action-2', action_type: 'purge_message', status: 'approved' },
      ],
    });

    const result = await remediateApprovedActions('campaign-1', 'operator@example.com');

    const updateCalls = callsContaining('UPDATE remediation_actions SET status');
    expect(updateCalls).toHaveLength(2);
    expect(writeAuditEventMock).toHaveBeenCalledTimes(2);
    expect(writeAuditEventMock.mock.calls[0][0]).toMatchObject({
      campaignId: 'campaign-1',
      eventType: 'remediation_completed',
    });
    expect(result.actions.every((a) => a.status === 'completed')).toBe(true);
  });

  it('is idempotent: a second call transitions nothing and writes no second audit row (REMED-04)', async () => {
    // First call: two approved rows.
    stage({
      remediationRows: [
        { id: 'action-1', action_type: 'block_sender', status: 'approved' },
        { id: 'action-2', action_type: 'purge_message', status: 'approved' },
      ],
    });
    await remediateApprovedActions('campaign-1', 'operator@example.com');
    expect(writeAuditEventMock).toHaveBeenCalledTimes(2);

    // Second call: same rows are now already completed — status filter finds nothing to transition.
    clientCalls = [];
    stage({
      remediationRows: [
        { id: 'action-1', action_type: 'block_sender', status: 'completed' },
        { id: 'action-2', action_type: 'purge_message', status: 'completed' },
      ],
    });
    const result2 = await remediateApprovedActions('campaign-1', 'operator@example.com');

    expect(callsContaining('UPDATE remediation_actions SET status')).toHaveLength(0);
    // Still exactly 2 total audit calls across both calls (no additional on the second).
    expect(writeAuditEventMock).toHaveBeenCalledTimes(2);
    expect(result2.actions.every((a) => a.alreadyCompleted)).toBe(true);
  });

  it('raises an explicit failure (RemediationValidationError) when the campaign has zero remediation_actions rows', async () => {
    stage({ remediationRows: [] });

    await expect(remediateApprovedActions('campaign-1', 'operator@example.com')).rejects.toThrow(
      RemediationValidationError
    );
    expect(writeAuditEventMock).not.toHaveBeenCalled();
  });

  it('calls generateAndPostAcknowledgment exactly once with the campaignId when an approved acknowledge_user row is remediated', async () => {
    stage({
      remediationRows: [{ id: 'action-1', action_type: 'acknowledge_user', status: 'approved' }],
    });

    await remediateApprovedActions('campaign-1', 'operator@example.com');

    expect(generateAndPostAcknowledgmentMock).toHaveBeenCalledTimes(1);
    expect(generateAndPostAcknowledgmentMock).toHaveBeenCalledWith('campaign-1');
  });

  it('does NOT call generateAndPostAcknowledgment for a block_sender/warn_user-only remediation', async () => {
    stage({
      remediationRows: [
        { id: 'action-1', action_type: 'block_sender', status: 'approved' },
        { id: 'action-2', action_type: 'warn_user', status: 'approved' },
      ],
    });

    await remediateApprovedActions('campaign-1', 'operator@example.com');

    expect(generateAndPostAcknowledgmentMock).not.toHaveBeenCalled();
  });

  it('does NOT call generateAndPostAcknowledgment when the acknowledge_user row staged is already completed (idempotent re-run)', async () => {
    stage({
      remediationRows: [{ id: 'action-1', action_type: 'acknowledge_user', status: 'completed' }],
    });

    await remediateApprovedActions('campaign-1', 'operator@example.com');

    expect(generateAndPostAcknowledgmentMock).not.toHaveBeenCalled();
  });

  it('does not propagate a generateAndPostAcknowledgment rejection out of remediateApprovedActions (already-committed transition)', async () => {
    stage({
      remediationRows: [{ id: 'action-1', action_type: 'acknowledge_user', status: 'approved' }],
    });
    generateAndPostAcknowledgmentMock.mockRejectedValueOnce(new Error('Autotask unavailable'));

    const result = await remediateApprovedActions('campaign-1', 'operator@example.com');

    expect(result.actions[0]).toMatchObject({ actionType: 'acknowledge_user', status: 'completed' });
  });
});

// =============================================================================
// markCampaignFalsePositive (D-04 guard)
// =============================================================================

describe('markCampaignFalsePositive', () => {
  it('throws RemediationConflictError when any remediation_actions row has status approved/completed', async () => {
    stage({
      guardRows: [{ id: 'action-1' }],
    });

    await expect(
      markCampaignFalsePositive('campaign-1', 'operator@example.com')
    ).rejects.toThrow(RemediationConflictError);
    expect(callsContaining('UPDATE campaigns SET status')).toHaveLength(0);
  });

  it('sets campaigns.status to false_positive and writes one audit row when there are no approved/completed rows', async () => {
    stage({
      guardRows: [],
      campaign: [{ status: 'open' }],
    });

    const result = await markCampaignFalsePositive('campaign-1', 'operator@example.com', 'confirmed benign');

    expect(callsContaining('UPDATE campaigns SET status')).toHaveLength(1);
    expect(writeAuditEventMock).toHaveBeenCalledTimes(1);
    const [auditArgs] = writeAuditEventMock.mock.calls[0];
    expect(auditArgs).toMatchObject({
      campaignId: 'campaign-1',
      eventType: 'campaign_marked_false_positive',
      payload: { previousStatus: 'open', reason: 'confirmed benign' },
    });
    expect(result).toMatchObject({ campaignId: 'campaign-1', status: 'false_positive' });
  });

  it('throws RemediationValidationError when the campaign does not exist', async () => {
    stage({
      guardRows: [],
      campaign: [],
    });

    await expect(
      markCampaignFalsePositive('campaign-missing', 'operator@example.com')
    ).rejects.toThrow(RemediationValidationError);
  });
});