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

import { describe, test, expect, beforeAll, afterAll, mock } from "bun:test";
import { db } from "@/db";
import {
  annotators,
  paragraphs,
  assignments,
  humanLabels,
  quizSessions,
} from "@/db/schema";
import { eq, and } 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 TEST_ANNOTATOR = {
  id: "test-label-user",
  displayName: "Test Label User",
  password: "testpass",
};

const TEST_PARAGRAPHS = [
  {
    id: "test-para-1",
    text: "The board oversees cybersecurity risk.",
    wordCount: 6,
    paragraphIndex: 0,
    companyName: "Test Corp",
    cik: "0001234567",
    ticker: "TEST",
    filingType: "10-K",
    filingDate: "2024-03-15",
    fiscalYear: 2024,
    accessionNumber: "0001234567-24-000001",
    secItem: "Item 1C",
  },
  {
    id: "test-para-2",
    text: "Our CISO manages the cybersecurity program.",
    wordCount: 7,
    paragraphIndex: 1,
    companyName: "Test Corp",
    cik: "0001234567",
    ticker: "TEST",
    filingType: "10-K",
    filingDate: "2024-03-15",
    fiscalYear: 2024,
    accessionNumber: "0001234567-24-000001",
    secItem: "Item 1C",
  },
];

const QUIZ_SESSION_ID = "test-label-quiz-session";

beforeAll(async () => {
  // Clean up in dependency order
  await db
    .delete(humanLabels)
    .where(eq(humanLabels.annotatorId, TEST_ANNOTATOR.id));
  await db
    .delete(assignments)
    .where(eq(assignments.annotatorId, TEST_ANNOTATOR.id));
  await db
    .delete(quizSessions)
    .where(eq(quizSessions.annotatorId, TEST_ANNOTATOR.id));
  for (const p of TEST_PARAGRAPHS) {
    await db.delete(paragraphs).where(eq(paragraphs.id, p.id));
  }
  await db.delete(annotators).where(eq(annotators.id, TEST_ANNOTATOR.id));

  // Create test data
  await db.insert(annotators).values(TEST_ANNOTATOR);
  await db.insert(paragraphs).values(TEST_PARAGRAPHS);
  await db.insert(assignments).values(
    TEST_PARAGRAPHS.map((p) => ({
      paragraphId: p.id,
      annotatorId: TEST_ANNOTATOR.id,
    })),
  );

  // Create a passed quiz session with warmup completed
  await db.insert(quizSessions).values({
    id: QUIZ_SESSION_ID,
    annotatorId: TEST_ANNOTATOR.id,
    totalQuestions: 8,
    score: 8,
    passed: true,
    completedAt: new Date(),
    answers: JSON.stringify({ quizAnswers: [], warmupCompleted: 5 }),
  });

  // Set up session cookie
  cookieJar.clear();
  await createSession(TEST_ANNOTATOR.id);
});

afterAll(async () => {
  await db
    .delete(humanLabels)
    .where(eq(humanLabels.annotatorId, TEST_ANNOTATOR.id));
  await db
    .delete(assignments)
    .where(eq(assignments.annotatorId, TEST_ANNOTATOR.id));
  await db
    .delete(quizSessions)
    .where(eq(quizSessions.annotatorId, TEST_ANNOTATOR.id));
  for (const p of TEST_PARAGRAPHS) {
    await db.delete(paragraphs).where(eq(paragraphs.id, p.id));
  }
  await db.delete(annotators).where(eq(annotators.id, TEST_ANNOTATOR.id));
});

describe("GET /api/label", () => {
  test("returns first assigned paragraph", async () => {
    const res = await GET();
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.paragraph).toBeDefined();
    expect(data.paragraph.id).toBe("test-para-1");
    expect(data.paragraph.text).toBe(TEST_PARAGRAPHS[0].text);
    expect(data.paragraph.companyName).toBe("Test Corp");
    expect(data.paragraph.filingType).toBe("10-K");
    expect(data.progress).toBeDefined();
    expect(data.progress.completed).toBe(0);
    expect(data.progress.total).toBe(2);
  });

  test("redirects to warmup if warmup not complete", async () => {
    // Temporarily modify quiz session to remove warmup completion
    await db
      .update(quizSessions)
      .set({ answers: "[]" })
      .where(eq(quizSessions.id, QUIZ_SESSION_ID));

    const res = await GET();
    const data = await res.json();
    expect(data.redirectToWarmup).toBe(true);

    // Restore warmup completion
    await db
      .update(quizSessions)
      .set({
        answers: JSON.stringify({ quizAnswers: [], warmupCompleted: 5 }),
      })
      .where(eq(quizSessions.id, QUIZ_SESSION_ID));
  });
});

describe("POST /api/label", () => {
  test("creates a label for assigned paragraph", async () => {
    const req = new Request("http://localhost:3000/api/label", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "test-para-1",
        contentCategory: "Board Governance",
        specificityLevel: 1,
        sessionId: "test-session-1",
        durationMs: 5000,
      }),
    });

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

    expect(res.status).toBe(200);
    expect(data.ok).toBe(true);
    expect(data.progress.completed).toBe(1);
    expect(data.progress.total).toBe(2);
  });

  test("GET returns next unlabeled paragraph after labeling", async () => {
    const res = await GET();
    const data = await res.json();

    expect(res.status).toBe(200);
    expect(data.paragraph.id).toBe("test-para-2");
    expect(data.progress.completed).toBe(1);
    expect(data.progress.total).toBe(2);
  });

  test("rejects duplicate label", async () => {
    const req = new Request("http://localhost:3000/api/label", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "test-para-1",
        contentCategory: "Board Governance",
        specificityLevel: 1,
        sessionId: "test-session-1",
        durationMs: 3000,
      }),
    });

    const res = await POST(req);
    expect(res.status).toBe(409);
    const data = await res.json();
    expect(data.error).toBe("Already labeled this paragraph");
  });

  test("rejects invalid category", async () => {
    const req = new Request("http://localhost:3000/api/label", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "test-para-2",
        contentCategory: "Invalid Category",
        specificityLevel: 1,
        sessionId: "test-session-1",
      }),
    });

    const res = await POST(req);
    expect(res.status).toBe(400);
    const data = await res.json();
    expect(data.error).toBe("Invalid content category");
  });

  test("rejects invalid specificity", async () => {
    const req = new Request("http://localhost:3000/api/label", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "test-para-2",
        contentCategory: "Board Governance",
        specificityLevel: 5,
        sessionId: "test-session-1",
      }),
    });

    const res = await POST(req);
    expect(res.status).toBe(400);
    const data = await res.json();
    expect(data.error).toBe("Invalid specificity level");
  });

  test("rejects unassigned paragraph", async () => {
    const req = new Request("http://localhost:3000/api/label", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "nonexistent-para",
        contentCategory: "Board Governance",
        specificityLevel: 1,
        sessionId: "test-session-1",
      }),
    });

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

  test("labels second paragraph and GET returns done", async () => {
    const req = new Request("http://localhost:3000/api/label", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        paragraphId: "test-para-2",
        contentCategory: "Management Role",
        specificityLevel: 2,
        notes: "Test note",
        sessionId: "test-session-1",
        durationMs: 4000,
      }),
    });

    const res = await POST(req);
    const data = await res.json();
    expect(res.status).toBe(200);
    expect(data.ok).toBe(true);
    expect(data.progress.completed).toBe(2);
    expect(data.progress.total).toBe(2);

    // GET should now return done
    const getRes = await GET();
    const getData = await getRes.json();
    expect(getData.done).toBe(true);
    expect(getData.progress.completed).toBe(2);
  });
});

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