Răsfoiți Sursa

test(ledger): add trigger verification test infrastructure

ding113 4 zile în urmă
părinte
comite
76c20cb392

+ 3 - 0
src/lib/ledger-backfill/index.ts

@@ -0,0 +1,3 @@
+import "server-only";
+
+export { backfillUsageLedger } from "./service";

+ 80 - 0
src/lib/ledger-backfill/service.ts

@@ -0,0 +1,80 @@
+import "server-only";
+
+import { sql } from "drizzle-orm";
+import { db } from "@/drizzle/db";
+
+export async function backfillUsageLedger(): Promise<{ inserted: number }> {
+  const LOCK_KEY = 20260101;
+
+  const result = await db.execute(sql`
+    SELECT pg_try_advisory_lock(${LOCK_KEY}) AS acquired
+  `);
+
+  const acquired = (result as unknown as Array<{ acquired: boolean }>)[0]?.acquired;
+  if (!acquired) {
+    return { inserted: 0 };
+  }
+
+  try {
+    const insertResult = await db.execute(sql`
+      INSERT INTO usage_ledger (
+        request_id, user_id, key, provider_id, final_provider_id,
+        model, original_model, endpoint, api_type, session_id,
+        status_code, is_success, blocked_by,
+        cost_usd, cost_multiplier,
+        input_tokens, output_tokens,
+        cache_creation_input_tokens, cache_read_input_tokens,
+        cache_creation_5m_input_tokens, cache_creation_1h_input_tokens,
+        cache_ttl_applied, context_1m_applied, swap_cache_ttl_applied,
+        duration_ms, ttfb_ms, created_at
+      )
+      SELECT
+        mr.id,
+        mr.user_id,
+        mr.key,
+        mr.provider_id,
+        COALESCE(
+          CASE
+            WHEN mr.provider_chain IS NOT NULL
+              AND jsonb_typeof(mr.provider_chain) = 'array'
+              AND jsonb_array_length(mr.provider_chain) > 0
+              AND jsonb_typeof(mr.provider_chain -> -1) = 'object'
+              AND (mr.provider_chain -> -1 ? 'id')
+              AND (mr.provider_chain -> -1 ->> 'id') ~ '^[0-9]+$'
+            THEN (mr.provider_chain -> -1 ->> 'id')::integer
+          END,
+          mr.provider_id
+        ),
+        mr.model,
+        mr.original_model,
+        mr.endpoint,
+        mr.api_type,
+        mr.session_id,
+        mr.status_code,
+        (mr.error_message IS NULL OR mr.error_message = ''),
+        mr.blocked_by,
+        mr.cost_usd,
+        mr.cost_multiplier,
+        mr.input_tokens,
+        mr.output_tokens,
+        mr.cache_creation_input_tokens,
+        mr.cache_read_input_tokens,
+        mr.cache_creation_5m_input_tokens,
+        mr.cache_creation_1h_input_tokens,
+        mr.cache_ttl_applied,
+        mr.context_1m_applied,
+        mr.swap_cache_ttl_applied,
+        mr.duration_ms,
+        mr.ttfb_ms,
+        mr.created_at
+      FROM message_request mr
+      WHERE mr.blocked_by IS DISTINCT FROM 'warmup'
+      ON CONFLICT (request_id) DO NOTHING
+    `);
+
+    const inserted = Number((insertResult as unknown as { rowCount?: number }).rowCount ?? 0);
+    return { inserted };
+  } finally {
+    await db.execute(sql`SELECT pg_advisory_unlock(${LOCK_KEY})`);
+  }
+}

+ 37 - 0
tests/unit/usage-ledger/backfill.test.ts

@@ -0,0 +1,37 @@
+import { readFileSync } from "node:fs";
+import { resolve } from "node:path";
+import { describe, expect, it, vi } from "vitest";
+
+process.env.DSN = "";
+
+vi.mock("@/drizzle/db", () => ({
+  db: {
+    execute: vi.fn(),
+  },
+}));
+
+vi.mock("drizzle-orm", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("drizzle-orm")>();
+  return { ...actual, sql: actual.sql };
+});
+
+const { backfillUsageLedger } = await import("@/lib/ledger-backfill");
+
+const serviceSource = readFileSync(
+  resolve(process.cwd(), "src/lib/ledger-backfill/service.ts"),
+  "utf-8"
+);
+
+describe("backfillUsageLedger", () => {
+  it("exports backfillUsageLedger function", () => {
+    expect(typeof backfillUsageLedger).toBe("function");
+  });
+
+  it("uses ON CONFLICT in backfill SQL", () => {
+    expect(serviceSource).toContain("ON CONFLICT");
+  });
+
+  it("uses DO NOTHING in backfill SQL", () => {
+    expect(serviceSource).toContain("DO NOTHING");
+  });
+});

+ 57 - 0
tests/unit/usage-ledger/repository.test.ts

@@ -0,0 +1,57 @@
+import { describe, expect, it, vi } from "vitest";
+
+process.env.DSN = "";
+
+vi.mock("@/drizzle/db", () => ({
+  db: {
+    select: vi.fn(() => ({ from: vi.fn(() => ({ where: vi.fn(async () => []) })) })),
+    execute: vi.fn(async () => []),
+  },
+}));
+
+vi.mock("drizzle-orm", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("drizzle-orm")>();
+  return {
+    ...actual,
+    and: vi.fn((...args: unknown[]) => args),
+    eq: vi.fn((...args: unknown[]) => args),
+    gte: vi.fn((...args: unknown[]) => args),
+    lt: vi.fn((...args: unknown[]) => args),
+    inArray: vi.fn((...args: unknown[]) => args),
+  };
+});
+
+vi.mock("@/drizzle/schema", () => ({
+  usageLedger: {
+    userId: "user_id",
+    key: "key",
+    finalProviderId: "final_provider_id",
+    costUsd: "cost_usd",
+    createdAt: "created_at",
+    blockedBy: "blocked_by",
+  },
+}));
+
+vi.mock("@/repository/_shared/ledger-conditions", () => ({
+  LEDGER_BILLING_CONDITION: {},
+}));
+
+const repo = await import("@/repository/usage-ledger");
+
+describe("usage-ledger repository", () => {
+  it("exports sumLedgerCostInTimeRange", () => {
+    expect(typeof repo.sumLedgerCostInTimeRange).toBe("function");
+  });
+
+  it("exports sumLedgerTotalCost", () => {
+    expect(typeof repo.sumLedgerTotalCost).toBe("function");
+  });
+
+  it("exports sumLedgerTotalCostBatch", () => {
+    expect(typeof repo.sumLedgerTotalCostBatch).toBe("function");
+  });
+
+  it("exports countLedgerRequestsInTimeRange", () => {
+    expect(typeof repo.countLedgerRequestsInTimeRange).toBe("function");
+  });
+});

+ 31 - 0
tests/unit/usage-ledger/trigger.test.ts

@@ -0,0 +1,31 @@
+import { readFileSync } from "node:fs";
+import { resolve } from "node:path";
+import { describe, expect, it } from "vitest";
+
+const sql = readFileSync(resolve(process.cwd(), "src/lib/ledger-backfill/trigger.sql"), "utf-8");
+
+describe("fn_upsert_usage_ledger trigger SQL", () => {
+  it("contains warmup exclusion check", () => {
+    expect(sql).toContain("blocked_by = 'warmup'");
+  });
+
+  it("contains ON CONFLICT UPSERT", () => {
+    expect(sql).toContain("ON CONFLICT (request_id) DO UPDATE");
+  });
+
+  it("contains EXCEPTION error handling", () => {
+    expect(sql).toContain("EXCEPTION WHEN OTHERS");
+  });
+
+  it("pre-validates provider_chain before extraction", () => {
+    expect(sql).toContain("jsonb_typeof");
+  });
+
+  it("computes is_success from error_message", () => {
+    expect(sql).toContain("error_message IS NULL");
+  });
+
+  it("creates trigger binding", () => {
+    expect(sql).toContain("CREATE TRIGGER trg_upsert_usage_ledger");
+  });
+});