UHD_Media_Tool/apps/server/test/auth-settings/harness.ts

import { type DB, migrateToLatest } from "@umt/db";
import { migrateEngineSchema } from "@umt/tasks";
import { Kysely, PostgresDialect } from "kysely";
import pg from "pg";
import { vi } from "vitest";
import { ConfigSchema } from "../../src/config.ts";
import { createOutboxPort } from "../../src/context.ts";
import { createApp } from "../../src/http/app.ts";
import { EventBus } from "../../src/http/sse.ts";
import { createCapabilityRegistry } from "../../src/modules/agent/capability-registry.ts";
import { createAuthSettingsModule } from "../../src/modules/auth-settings/index.ts";
import { createSecretStore, createSettingsReader } from "../../src/modules/auth-settings/provided.ts";
import type { ServerContext } from "../../src/modules/types.ts";
import { DEV_KEY, memoryLogger } from "../helpers.ts";

// 真实 PG 的 W2-A 测试夹具:独立业务 schema + 独立引擎 schema;任务引擎用假对象(只测 cancel 透传)。

export const PG_URL = process.env.DATABASE_URL ?? "postgres://umt:umt@localhost:55433/umt";
export const MASTER_KEY = "test-master-key-0123456789";

export async function pgAvailable() {
  if (process.env.UMT_SKIP_DB_TESTS === "1") return false;
  const c = new pg.Client({ connectionString: PG_URL, connectionTimeoutMillis: 2000 });
  try {
    await c.connect();
    await c.end();
    return true;
  } catch {
    return false;
  }
}

/**
 * Kysely 迁移器会反射库里所有 schema 的表;别的测试文件同时 drop schema 时反射会报"schema 不存在"。
 * 同一进程内的测试文件用 PG 咨询锁把"迁移 / 删 schema"串行化,另加重试兜底(其它进程的测试也在用同一个库)。
 */
const SCHEMA_LOCK = 0x756d7431;
export async function withSchemaLock<T>(fn: () => Promise<T>): Promise<T> {
  const c = new pg.Client({ connectionString: PG_URL });
  await c.connect();
  try {
    await c.query("select pg_advisory_lock($1)", [SCHEMA_LOCK]);
    for (let attempt = 1; ; attempt++) {
      try {
        return await fn();
      } catch (err) {
        if (attempt >= 4 || !/does not exist/.test(String((err as Error)?.message))) throw err;
        await new Promise((r) => setTimeout(r, 200 * attempt));
      }
    }
  } finally {
    await c.query("select pg_advisory_unlock($1)", [SCHEMA_LOCK]).catch(() => {});
    await c.end();
  }
}

async function adminQuery(q: string) {
  const c = new pg.Client({ connectionString: PG_URL });
  await c.connect();
  try {
    await c.query(q);
  } finally {
    await c.end();
  }
}

export async function createHarness(opts: { engine?: boolean } = {}) {
  const uniq = `${process.pid.toString(36)}${Math.random().toString(36).slice(2, 7)}`;
  const schema = `w2a_test_${uniq}`;
  const engineSchema = `${schema}_ow`;
  await adminQuery(`create schema ${schema}`);
  const db = new Kysely<DB>({
    dialect: new PostgresDialect({
      pool: new pg.Pool({ connectionString: PG_URL, max: 5, options: `-c search_path=${schema},public` }),
    }),
  });
  await withSchemaLock(async () => {
    await migrateToLatest(db, { migrationTableSchema: schema });
    if (opts.engine) await migrateEngineSchema(PG_URL, engineSchema);
  });
  const { logger, lines } = memoryLogger();
  const bus = new EventBus(100);
  const config = ConfigSchema.parse({
    DATABASE_URL: PG_URL,
    SECRET_MASTER_KEY: MASTER_KEY,
    NODE_ENV: "development",
    DEV_ADMIN_KEY: DEV_KEY,
  });
  const cancel = vi.fn(async (_runId: string, _reason?: string) => true);
  const ctx = {
    role: "all",
    db,
    config,
    logger,
    events: bus,
    engine: { cancel } as unknown as ServerContext["engine"],
    outbox: createOutboxPort(),
    secrets: createSecretStore({ db, masterKey: MASTER_KEY }),
    settings: createSettingsReader({ db, logger, events: bus }),
    capabilities: createCapabilityRegistry({ logger }),
  } as ServerContext;
  const mod = createAuthSettingsModule({ setupTokenOnStart: false, engineSchema });
  const app = createApp({ config, logger, bus, db, role: "all", version: "0.0.0-test", modules: [mod], ctx });
  return {
    db,
    app,
    ctx,
    bus,
    lines,
    cancel,
    schema,
    engineSchema,
    async close() {
      await db.destroy();
      await withSchemaLock(async () => {
        await adminQuery(`drop schema if exists ${schema} cascade`);
        await adminQuery(`drop schema if exists ${engineSchema} cascade`);
      });
    },
  };
}
export type Harness = Awaited<ReturnType<typeof createHarness>>;

/** 极简 cookie jar:记住 Set-Cookie,请求时带上;写方法自动带 CSRF 头 */
export class Browser {
  cookies = new Map<string, string>();
  constructor(
    private readonly app: Harness["app"],
    readonly ip = "203.0.113.10",
  ) {}

  absorb(res: Response) {
    for (const sc of res.headers.getSetCookie()) {
      const [pair] = sc.split(";");
      const i = pair!.indexOf("=");
      const k = pair!.slice(0, i).trim();
      const v = pair!.slice(i + 1).trim();
      if (!v || /max-age=0/i.test(sc)) this.cookies.delete(k);
      else this.cookies.set(k, decodeURIComponent(v));
    }
  }

  async req(
    method: string,
    path: string,
    body?: unknown,
    opts: { csrf?: boolean; headers?: Record<string, string> } = {},
  ) {
    const headers: Record<string, string> = { "x-real-ip": this.ip, ...opts.headers };
    if (this.cookies.size)
      headers.cookie = [...this.cookies].map(([k, v]) => `${k}=${encodeURIComponent(v)}`).join("; ");
    if (body !== undefined) headers["content-type"] = "application/json";
    const csrf = this.cookies.get("umt_csrf");
    if (method !== "GET" && opts.csrf !== false && csrf) headers["x-csrf-token"] = csrf;
    const res = await this.app.request(`/api/v1${path}`, {
      method,
      headers,
      ...(body !== undefined ? { body: JSON.stringify(body) } : {}),
    });
    this.absorb(res);
    return res;
  }
}

/** Bearer 调用 */
export function bearer(app: Harness["app"], key: string) {
  return (method: string, path: string, body?: unknown) =>
    app.request(`/api/v1${path}`, {
      method,
      headers: {
        authorization: `Bearer ${key}`,
        ...(body !== undefined ? { "content-type": "application/json" } : {}),
      },
      ...(body !== undefined ? { body: JSON.stringify(body) } : {}),
    });
}