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(fn: () => Promise): Promise { 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({ 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>; /** 极简 cookie jar:记住 Set-Cookie,请求时带上;写方法自动带 CSRF 头 */ export class Browser { cookies = new Map(); 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 } = {}, ) { const headers: Record = { "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) } : {}), }); }