SEC-cyBERT/labelapp/app/api/adjudicate/__test__/adjudicate.test.ts

import { describe, test, expect, beforeAll, afterAll, mock } from "bun:test";
import { db } from "@/db";
import {
  annotators,
  paragraphs,
  assignments,
  humanLabels,
  adjudications,
} from "@/db/schema";
import { eq, and, inArray } from "drizzle-orm";

const cookieJar = new Map<string, string>();
const mockCookieStore = {
  get(name: string) {
    const value = cookieJar.get(name);
    return value ? { value } : undefined;
  },
  set(name: string, value: string, _opts?: unknown) {
    cookieJar.set(name, value);
  },
  delete(name: string) {
    cookieJar.delete(name);
  },
};

mock.module("next/headers", () => ({
  cookies: async () => mockCookieStore,
}));

const { createSession } = await import("@/lib/auth");
const { GET, POST } = await import("../route");

const ADMIN_USER = {
  id: "joey",
  displayName: "Joey",
  password: "adminpass",
};

const TEST_ANNOTATORS = [
  { id: "adj-ann-1", displayName: "Annotator 1", password: "pass1" },
  { id: "adj-ann-2", displayName: "Annotator 2", password: "pass2" },
  { id: "adj-ann-3", displayName: "Annotator 3", password: "pass3" },
];

const TEST_PARAGRAPHS = [
  {
    id: "adj-para-agree",
    text: "The board of directors oversees the company cybersecurity risk management.",
    wordCount: 11,
    paragraphIndex: 0,
    companyName: "Test Corp",
    cik: "0009999999",
    ticker: "ADJT",
    filingType: "10-K",
    filingDate: "2024-03-15",
    fiscalYear: 2024,
    accessionNumber: "0009999999-24-000001",
    secItem: "Item 1C",
  },
  {
    id: "adj-para-disagree",
    text: "We employ third-party vendors and manage internal risk processes simultaneously.",
    wordCount: 10,
    paragraphIndex: 1,
    companyName: "Test Corp",
    cik: "0009999999",
    ticker: "ADJT",
    filingType: "10-K",
    filingDate: "2024-03-15",
    fiscalYear: 2024,
    accessionNumber: "0009999999-24-000001",
    secItem: "Item 1C",
  },
  {
    id: "adj-para-majority",
    text: "Our CISO reports quarterly to the board on cybersecurity matters.",
    wordCount: 10,
    paragraphIndex: 2,
    companyName: "Test Corp",
    cik: "0009999999",
    ticker: "ADJT",
    filingType: "10-K",
    filingDate: "2024-03-15",
    fiscalYear: 2024,
    accessionNumber: "0009999999-24-000001",
    secItem: "Item 1C",
  },
];

const PARAGRAPH_IDS = TEST_PARAGRAPHS.map((p) => p.id);
const ALL_ANNOTATOR_IDS = [
  ADMIN_USER.id,
  ...TEST_ANNOTATORS.map((a) => a.id),
];

async function cleanup() {
  // Clean in dependency order
  await db
    .delete(adjudications)
    .where(inArray(adjudications.paragraphId, PARAGRAPH_IDS));
  await db
    .delete(humanLabels)
    .where(inArray(humanLabels.paragraphId, PARAGRAPH_IDS));
  await db
    .delete(assignments)
    .where(inArray(assignments.paragraphId, PARAGRAPH_IDS));
  await db.delete(paragraphs).where(inArray(paragraphs.id, PARAGRAPH_IDS));
  for (const a of TEST_ANNOTATORS) {
    await db.delete(annotators).where(eq(annotators.id, a.id));
  }
  // Only delete admin if it exists (don't fail if not present)
  await db.delete(annotators).where(eq(annotators.id, ADMIN_USER.id));
}

beforeAll(async () => {
  await cleanup();

  // Create test data
  await db
    .insert(annotators)
    .values([ADMIN_USER, ...TEST_ANNOTATORS])
    .onConflictDoNothing();
  await db.insert(paragraphs).values(TEST_PARAGRAPHS);

  // Assign all paragraphs to all 3 annotators
  const assignmentRows = TEST_PARAGRAPHS.flatMap((p) =>
    TEST_ANNOTATORS.map((a) => ({
      paragraphId: p.id,
      annotatorId: a.id,
    })),
  );
  await db.insert(assignments).values(assignmentRows);

  // Set up admin session
  cookieJar.clear();
  await createSession(ADMIN_USER.id);
});

afterAll(async () => {
  await cleanup();
});

describe("GET /api/adjudicate", () => {
  test("returns empty queue when no labels exist", async () => {
    const res = await GET();
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.queue).toEqual([]);
  });

  test("auto-resolves consensus (3/3 agree) and returns empty queue", async () => {
    // Insert 3 agreeing labels for adj-para-agree
    const labels = TEST_ANNOTATORS.map((a) => ({
      paragraphId: "adj-para-agree",
      annotatorId: a.id,
      contentCategory: "Board Governance",
      specificityLevel: 3,
      sessionId: "test-adj-session",
    }));
    await db.insert(humanLabels).values(labels);

    const res = await GET();
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.queue).toEqual([]);

    // Verify adjudication was created
    const [adj] = await db
      .select()
      .from(adjudications)
      .where(eq(adjudications.paragraphId, "adj-para-agree"));

    expect(adj).toBeDefined();
    expect(adj.method).toBe("consensus");
    expect(adj.finalCategory).toBe("Board Governance");
    expect(adj.finalSpecificity).toBe(3);
  });

  test("auto-resolves majority (2/3 agree) and returns empty queue", async () => {
    // Insert 2 agreeing + 1 disagreeing for adj-para-majority
    await db.insert(humanLabels).values([
      {
        paragraphId: "adj-para-majority",
        annotatorId: TEST_ANNOTATORS[0].id,
        contentCategory: "Management Role",
        specificityLevel: 2,
        sessionId: "test-adj-session",
      },
      {
        paragraphId: "adj-para-majority",
        annotatorId: TEST_ANNOTATORS[1].id,
        contentCategory: "Management Role",
        specificityLevel: 2,
        sessionId: "test-adj-session",
      },
      {
        paragraphId: "adj-para-majority",
        annotatorId: TEST_ANNOTATORS[2].id,
        contentCategory: "Board Governance",
        specificityLevel: 3,
        sessionId: "test-adj-session",
      },
    ]);

    const res = await GET();
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.queue).toEqual([]);

    // Verify adjudication was created
    const [adj] = await db
      .select()
      .from(adjudications)
      .where(eq(adjudications.paragraphId, "adj-para-majority"));

    expect(adj).toBeDefined();
    expect(adj.method).toBe("majority");
    expect(adj.finalCategory).toBe("Management Role");
    expect(adj.finalSpecificity).toBe(2);
  });

  test("returns 3-way disagreement in queue", async () => {
    // Insert 3 fully disagreeing labels for adj-para-disagree
    await db.insert(humanLabels).values([
      {
        paragraphId: "adj-para-disagree",
        annotatorId: TEST_ANNOTATORS[0].id,
        contentCategory: "Third-Party Risk",
        specificityLevel: 1,
        sessionId: "test-adj-session",
      },
      {
        paragraphId: "adj-para-disagree",
        annotatorId: TEST_ANNOTATORS[1].id,
        contentCategory: "Risk Management Process",
        specificityLevel: 2,
        sessionId: "test-adj-session",
      },
      {
        paragraphId: "adj-para-disagree",
        annotatorId: TEST_ANNOTATORS[2].id,
        contentCategory: "Management Role",
        specificityLevel: 3,
        sessionId: "test-adj-session",
      },
    ]);

    const res = await GET();
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.queue).toHaveLength(1);
    expect(data.queue[0].paragraphId).toBe("adj-para-disagree");
    expect(data.queue[0].labels).toHaveLength(3);
    expect(data.queue[0].splitSeverity).toBe(3);
  });
});

describe("POST /api/adjudicate", () => {
  test("resolves an adjudication", async () => {
    const req = new Request("http://localhost:3000/api/adjudicate", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "adj-para-disagree",
        finalCategory: "Third-Party Risk",
        finalSpecificity: 2,
        notes: "Resolved via discussion",
      }),
    });

    const res = await POST(req);
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.ok).toBe(true);

    // Verify adjudication row
    const [adj] = await db
      .select()
      .from(adjudications)
      .where(eq(adjudications.paragraphId, "adj-para-disagree"));

    expect(adj).toBeDefined();
    expect(adj.method).toBe("discussion");
    expect(adj.finalCategory).toBe("Third-Party Risk");
    expect(adj.finalSpecificity).toBe(2);
    expect(adj.adjudicatorId).toBe("joey");
    expect(adj.notes).toBe("Resolved via discussion");
  });

  test("GET returns empty queue after resolution", async () => {
    const res = await GET();
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.queue).toEqual([]);
  });

  test("rejects invalid category", async () => {
    const req = new Request("http://localhost:3000/api/adjudicate", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "adj-para-disagree",
        finalCategory: "Invalid",
        finalSpecificity: 2,
      }),
    });

    const res = await POST(req);
    expect(res.status).toBe(400);
  });

  test("rejects invalid specificity", async () => {
    const req = new Request("http://localhost:3000/api/adjudicate", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "adj-para-disagree",
        finalCategory: "Board Governance",
        finalSpecificity: 5,
      }),
    });

    const res = await POST(req);
    expect(res.status).toBe(400);
  });
});

describe("GET /api/adjudicate - unauthorized", () => {
  test("returns 401 without session", async () => {
    cookieJar.clear();
    const res = await GET();
    expect(res.status).toBe(401);
  });

  test("returns 403 for non-admin user", async () => {
    cookieJar.clear();
    await createSession(TEST_ANNOTATORS[0].id);
    const res = await GET();
    expect(res.status).toBe(403);

    // Restore admin session for other tests
    cookieJar.clear();
    await createSession(ADMIN_USER.id);
  });
});