cfw-attachment/tests/fakes.ts

import { DatabaseSync } from "node:sqlite";
import migrationSql from "../migrations/0001_initial.sql?raw";
import migration2Sql from "../migrations/0002_upload_sessions.sql?raw";
import migration3Sql from "../migrations/0003_attachment_metadata.sql?raw";
import { vi } from "vitest";
import type { Env, FetcherLike } from "../src/env";

/**
 * 内存版 D1,基于 node:sqlite 执行真实 SQL。
 * 直接复用 migrations/0001 + 0002 建表,保证 store 层 SQL 被真实跑过。
 */
export class FakeD1 {
  private db: DatabaseSync;

  constructor() {
    this.db = new DatabaseSync(":memory:");
    this.db.exec(migrationSql);
    this.db.exec(migration2Sql);
    this.db.exec(migration3Sql);
  }

  prepare(sql: string) {
    const exec = (values: unknown[]) => {
      const statement = this.db.prepare(sql);
      return {
        first: async <T = unknown>(): Promise<T | null> => (statement.get(...values) as T | null) ?? null,
        all: async <T = unknown>() => ({ results: statement.all(...values) as T[], success: true as const }),
        run: async <T = unknown>() => {
          statement.run(...values);
          return { results: [] as T[], success: true as const };
        },
        bind: (...vals: unknown[]) => exec(vals),
      };
    };
    return exec([]);
  }
}

function fakeEtag(bytes: Uint8Array): string {
  let hash = 5381;
  for (let i = 0; i < bytes.length; i += 1) {
    hash = ((hash << 5) + hash + bytes[i]) >>> 0;
  }
  return hash.toString(16).padStart(8, "0");
}

interface StoredObject {
  body: Uint8Array;
  httpMetadata: R2HTTPMetadata;
  uploaded: number;
}

function toR2Object(key: string, stored: StoredObject): R2Object {
  return {
    key,
    size: stored.body.byteLength,
    etag: fakeEtag(stored.body),
    version: "fake-version",
    httpEtag: `"${fakeEtag(stored.body)}"`,
    uploaded: new Date(stored.uploaded),
    httpMetadata: stored.httpMetadata,
    customMetadata: {},
    storageClass: "Standard",
    writeHttpMetadata: () => new Headers(),
    storageUsage: 0,
  } as unknown as R2Object;
}

async function materializeBody(
  body: ReadableStream<Uint8Array> | ArrayBuffer | ArrayBufferView | string | null,
): Promise<Uint8Array> {
  if (body === null) {
    return new Uint8Array();
  }
  if (typeof body === "string") {
    return new TextEncoder().encode(body);
  }
  if (body instanceof ArrayBuffer) {
    return new Uint8Array(body);
  }
  if (ArrayBuffer.isView(body)) {
    return new Uint8Array(body.buffer, body.byteOffset, body.byteLength);
  }
  const buffer = await new Response(body as ReadableStream<Uint8Array>).arrayBuffer();
  return new Uint8Array(buffer);
}

interface MultipartSession {
  key: string;
  parts: Map<number, { bytes: Uint8Array; etag: string }>;
  httpMetadata: R2HTTPMetadata;
}

/**
 * 内存版 R2 bucket,用 Map 保存对象字节,覆盖 put/get/head/delete。
 */
export class FakeR2 {
  private store = new Map<string, StoredObject>();
  private sessions = new Map<string, MultipartSession>();
  private sessionCounter = 0;

  async put(
    key: string,
    body: ReadableStream<Uint8Array> | ArrayBuffer | ArrayBufferView | string | null,
    options: { httpMetadata?: R2HTTPMetadata } = {},
  ): Promise<R2Object> {
    const bytes = await materializeBody(body);
    const stored: StoredObject = { body: bytes, httpMetadata: options.httpMetadata ?? {}, uploaded: Date.now() };
    this.store.set(key, stored);
    return toR2Object(key, stored);
  }

  // --- 分片上传(内存版,镜像 R2 multipart 行为) ---

  async createMultipartUpload(
    key: string,
    options: { httpMetadata?: R2HTTPMetadata; customMetadata?: Record<string, string> } = {},
  ): Promise<R2MultipartUpload> {
    const uploadId = `mp-${++this.sessionCounter}`;
    this.sessions.set(uploadId, { key, parts: new Map(), httpMetadata: options.httpMetadata ?? {} });
    return this.multipartHandle(uploadId);
  }

  resumeMultipartUpload(key: string, uploadId: string): R2MultipartUpload {
    const session = this.sessions.get(uploadId);
    if (!session || session.key !== key) {
      throw new Error(`FakeR2: unknown multipart upload ${uploadId}`);
    }
    return this.multipartHandle(uploadId);
  }

  private multipartHandle(uploadId: string): R2MultipartUpload {
    const session = this.sessions.get(uploadId)!;
    const handle = {
      key: session.key,
      uploadId,
      uploadPart: async (
        partNumber: number,
        value: ReadableStream | ArrayBuffer | ArrayBufferView | string | Blob,
      ): Promise<R2UploadedPart> => {
        const bytes = await materializeBody(value as ReadableStream<Uint8Array>);
        const etag = fakeEtag(bytes);
        session.parts.set(partNumber, { bytes, etag });
        return { partNumber, etag } as R2UploadedPart;
      },
      complete: async (uploadedParts: R2UploadedPart[]): Promise<R2Object> => {
        const total = uploadedParts
          .slice()
          .sort((a, b) => a.partNumber - b.partNumber)
          .flatMap((part) => session.parts.get(part.partNumber)?.bytes ?? new Uint8Array());
        const merged = new Uint8Array(total.reduce((sum, part) => sum + part.length, 0));
        let offset = 0;
        for (const part of total) {
          merged.set(part, offset);
          offset += part.length;
        }
        const stored: StoredObject = {
          body: merged,
          httpMetadata: session.httpMetadata,
          uploaded: Date.now(),
        };
        this.store.set(session.key, stored);
        this.sessions.delete(uploadId);
        return toR2Object(session.key, stored);
      },
      abort: async (): Promise<void> => {
        this.sessions.delete(uploadId);
      },
    };
    return handle as unknown as R2MultipartUpload;
  }

  async get(key: string, _options?: unknown): Promise<R2ObjectBody | null> {
    const stored = this.store.get(key);
    if (!stored) return null;
    const head = toR2Object(key, stored);
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(stored.body);
        controller.close();
      },
    });
    return {
      ...head,
      body,
      bodyUsed: false,
      arrayBuffer: async () =>
        stored.body.buffer.slice(stored.body.byteOffset, stored.body.byteOffset + stored.body.byteLength),
      text: async () => new TextDecoder().decode(stored.body),
      json: async <T>() => JSON.parse(new TextDecoder().decode(stored.body)) as T,
      blob: async () => new Blob([stored.body as BlobPart]),
    } as unknown as R2ObjectBody;
  }

  async head(key: string): Promise<R2Object | null> {
    const stored = this.store.get(key);
    return stored ? toR2Object(key, stored) : null;
  }

  async delete(key: string): Promise<void> {
    this.store.delete(key);
  }

  has(key: string): boolean {
    return this.store.has(key);
  }
}

export function fakeAuth(userId: string | null): FetcherLike {
  return {
    fetch: vi.fn(async () =>
      Response.json(userId ? { user: { id: userId } } : null, { status: userId ? 200 : 401 }),
    ),
  };
}

export interface SharedStores {
  db: FakeD1;
  r2: FakeR2;
}

export function sharedStores(): SharedStores {
  return { db: new FakeD1(), r2: new FakeR2() };
}

export function makeEnv(
  stores: SharedStores,
  userId: string | null,
  maxUploadBytes = "52428800",
  publicBaseUrl = "https://attachment.local",
): Env {
  return {
    AUTH: fakeAuth(userId),
    DB: stores.db as unknown as D1Database,
    ATTACHMENTS: stores.r2 as unknown as R2Bucket,
    MAX_UPLOAD_BYTES: maxUploadBytes,
    PUBLIC_BASE_URL: publicBaseUrl,
  };
}