next-ai-draw-io/tests/unit/validate-diagram-route.test.ts

// @vitest-environment node
import { MockLanguageModelV3 } from "ai/test"
import { afterEach, describe, expect, it, vi } from "vitest"
import { POST as validateDiagram } from "@/app/api/validate-diagram/route"

const RESULT = {
    valid: false,
    issues: [
        {
            type: "overlap",
            severity: "critical",
            description: "Box A covers box B",
        },
    ],
    suggestions: ["Move box B to the right"],
}

// A vision model that answers with the result as JSON text
vi.mock("@/lib/ai-providers", () => ({
    getValidationModel: () =>
        new MockLanguageModelV3({
            doGenerate: (async () => ({
                content: [{ type: "text", text: JSON.stringify(RESULT) }],
                finishReason: { unified: "stop", raw: "stop" },
                usage: {
                    inputTokens: { total: 1 },
                    outputTokens: { total: 1 },
                },
                warnings: [],
            })) as any,
        }),
}))

// The quota, off unless a test turns it on
const quota = vi.hoisted(() => ({
    enabled: false,
    allowed: true,
    checks: [] as Array<{ limits: any; increment?: number }>,
    recorded: [] as number[],
}))
vi.mock("@/lib/dynamo-quota-manager", () => ({
    isQuotaEnabled: () => quota.enabled,
    checkAndIncrementRequest: async (
        _ip: string,
        limits: unknown,
        increment?: number,
    ) => {
        quota.checks.push({ limits, increment })
        return quota.allowed
            ? { allowed: true }
            : {
                  allowed: false,
                  type: "token",
                  error: "Daily token limit exceeded",
                  used: 10,
                  limit: 10,
              }
    },
    recordTokenUsage: async (_ip: string, tokens: number) => {
        quota.recorded.push(tokens)
    },
}))

const post = () =>
    new Request("http://localhost/api/validate-diagram", {
        method: "POST",
        headers: {
            "Content-Type": "application/json",
            "x-forwarded-for": "203.0.113.7",
        },
        body: JSON.stringify({ imageData: "data:image/png;base64,AAAA" }),
    })

afterEach(() => {
    delete process.env.ENABLE_VLM_VALIDATION
    quota.enabled = false
    quota.allowed = true
    quota.checks = []
    quota.recorded = []
})

describe("the quota", () => {
    it("refuses a check once the daily tokens are used up", async () => {
        quota.enabled = true
        quota.allowed = false
        const res = await validateDiagram(post())
        expect(res.status).toBe(429)
    })

    it("applies the token limits only, and records the tokens", async () => {
        // The request limit is for chats: the day's last chat must still
        // get its check, and the check does not count as a chat
        quota.enabled = true
        const res = await validateDiagram(post())
        expect(await res.json()).toEqual(RESULT)
        expect(quota.checks).toHaveLength(1)
        expect(quota.checks[0].increment).toBe(0)
        expect(quota.checks[0].limits.requests).toBe(0)
        expect(quota.recorded).toEqual([2])
    })
})

describe("POST /api/validate-diagram", () => {
    it("answers the model's result as JSON", async () => {
        const res = await validateDiagram(post())
        expect(res.headers.get("content-type")).toContain("application/json")
        expect(await res.json()).toEqual(RESULT)
    })

    it("answers valid when the check is turned off", async () => {
        process.env.ENABLE_VLM_VALIDATION = "false"
        const res = await validateDiagram(post())
        expect(await res.json()).toEqual({
            valid: true,
            issues: [],
            suggestions: [],
        })
    })
})