Shade/packages/shade-files/tests/integration/helpers/rig.ts

import { createShade, type Shade } from '@shade/sdk';
import {
  createPrekeyServer,
  MemoryPrekeyStore,
  PrekeyServerEvents,
} from '@shade/server';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import {
  PendingRpcRegistry,
  ShadeFileRpcChannel,
  attachClientRouting,
  attachFileHandler,
  createClientStreamsBridge,
  createFileClient,
  createFileHandler,
  createServerStreamsBridge,
  type ClientStreamsBridge,
  type FileClient,
  type FileHandler,
  type FileHandlerConfig,
  type ServerStreamsBridge,
} from '../../../src/index.js';

const crypto = new SubtleCryptoProvider();

export interface FileTestRig {
  alice: Shade;
  bob: Shade;
  fs: FileClient;
  bobHandler: FileHandler;
  /** Server-side streams bridge (Bob). */
  bobStreamsBridge: ServerStreamsBridge;
  /** Client-side streams bridge (Alice). */
  aliceStreamsBridge: ClientStreamsBridge;
  /** Tear everything down (kills servers, shuts down shades). */
  teardown(): Promise<void>;
}

/**
 * Setup options.
 *
 * Defaults to wiring streams-bridges on both sides so content I/O tests
 * (`read-write-streams.test.ts`) work transparently. Pass `withStreams: false`
 * to skip — useful for the legacy `std-ops` tests that don't need them.
 */
export interface SetupRigOptions {
  withStreams?: boolean;
}

export async function setupFileRig(
  bobConfig: Omit<FileHandlerConfig, 'streamsBridge'>,
  options: SetupRigOptions = {},
): Promise<FileTestRig> {
  const withStreams = options.withStreams ?? true;
  // 1. Prekey server
  const prekeyEvents = new PrekeyServerEvents();
  const prekey = createPrekeyServer({
    crypto,
    store: new MemoryPrekeyStore(),
    disableRateLimit: true,
    events: prekeyEvents,
  });
  const prekeyServer = Bun.serve({ port: 0, fetch: prekey.fetch });
  const prekeyUrl = `http://localhost:${prekeyServer.port}`;

  // 2. Two Shades
  const alice = await createShade({ prekeyServer: prekeyUrl, address: 'alice' });
  const bob = await createShade({ prekeyServer: prekeyUrl, address: 'bob' });

  // 3. BOTH sides need a transferRoute mounted because file-RPC is
  //    request/response — Bob's reply must reach Alice's HTTP endpoint
  //    just as Alice's request reached Bob's.
  let bobBaseUrl = '';
  let aliceBaseUrl = '';

  alice.configureTransfers({
    resolveBaseUrl: async (peer) => {
      if (peer === 'bob') return bobBaseUrl;
      throw new Error(`alice: unknown peer ${peer}`);
    },
  });
  bob.configureTransfers({
    resolveBaseUrl: async (peer) => {
      if (peer === 'alice') return aliceBaseUrl;
      throw new Error(`bob: unknown peer ${peer}`);
    },
  });

  const bobApp = await bob.transferRoute();
  const aliceApp = await alice.transferRoute();
  const bobServer = Bun.serve({ port: 0, fetch: bobApp.fetch });
  const aliceServer = Bun.serve({ port: 0, fetch: aliceApp.fetch });
  bobBaseUrl = `http://localhost:${bobServer.port}`;
  aliceBaseUrl = `http://localhost:${aliceServer.port}`;

  // 4. Streams bridges (both sides) — required for content I/O > 256 KiB.
  let bobStreamsBridge: ServerStreamsBridge | undefined;
  let aliceStreamsBridge: ClientStreamsBridge | undefined;
  if (withStreams) {
    bobStreamsBridge = await createServerStreamsBridge(bob);
    aliceStreamsBridge = await createClientStreamsBridge(alice);
  }

  // 5. Bob: file handler + channel
  const bobChannel = new ShadeFileRpcChannel(bob);
  const fullBobConfig: FileHandlerConfig = {
    ...bobConfig,
    ...(bobStreamsBridge !== undefined ? { streamsBridge: bobStreamsBridge } : {}),
  };
  const bobHandler = createFileHandler(bob, fullBobConfig);
  attachFileHandler(bobChannel, bobHandler);

  // 6. Alice: client + channel + pending registry
  const aliceChannel = new ShadeFileRpcChannel(alice);
  const alicePending = new PendingRpcRegistry();
  attachClientRouting(aliceChannel, alicePending);
  const fs = createFileClient(alice, aliceChannel, alicePending, 'bob', {
    defaultTimeoutMs: 5000,
    ...(aliceStreamsBridge !== undefined ? { streamsBridge: aliceStreamsBridge } : {}),
  });

  return {
    alice,
    bob,
    fs,
    bobHandler,
    bobStreamsBridge: bobStreamsBridge as ServerStreamsBridge,
    aliceStreamsBridge: aliceStreamsBridge as ClientStreamsBridge,
    async teardown() {
      bobChannel.destroy();
      aliceChannel.destroy();
      bobHandler.destroy();
      if (bobStreamsBridge !== undefined) await bobStreamsBridge.destroy();
      if (aliceStreamsBridge !== undefined) await aliceStreamsBridge.destroy();
      await alice.shutdown();
      await bob.shutdown();
      bobServer.stop();
      aliceServer.stop();
      prekeyServer.stop();
    },
  };
}