arboretum/packages/vscode/test/ws-client.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { BINARY_FRAME, FLOW, encodeBinaryFrame, type ClientMessage } from '@arboretum/shared';
import { ArbWsClient, type AttachmentSink, type RawSocket } from '../src/api/ws-client.js';

class FakeSocket implements RawSocket {
  sent: string[] = [];
  private openCb: (() => void) | null = null;
  private textCb: ((t: string) => void) | null = null;
  private binaryCb: ((d: Uint8Array) => void) | null = null;
  private closeCb: (() => void) | null = null;

  send(data: string): void {
    this.sent.push(data);
  }
  close(): void {
    this.closeCb?.();
  }
  onOpen(cb: () => void): void {
    this.openCb = cb;
  }
  onText(cb: (t: string) => void): void {
    this.textCb = cb;
  }
  onBinary(cb: (d: Uint8Array) => void): void {
    this.binaryCb = cb;
  }
  onClose(cb: () => void): void {
    this.closeCb = cb;
  }
  onError(): void {
    /* non utilisé dans ces tests */
  }

  fireOpen(): void {
    this.openCb?.();
  }
  fireText(msg: unknown): void {
    this.textCb?.(JSON.stringify(msg));
  }
  fireBinary(data: Uint8Array): void {
    this.binaryCb?.(data);
  }

  /** messages de contrôle envoyés, parsés et filtrés par type. */
  controls<T extends ClientMessage['type']>(type: T): Array<Extract<ClientMessage, { type: T }>> {
    return this.sent
      .map((s) => JSON.parse(s) as ClientMessage)
      .filter((m): m is Extract<ClientMessage, { type: T }> => m.type === type);
  }
}

function makeNoopSink(): AttachmentSink {
  return { data: () => {}, reset: () => {}, detached: () => {}, controlChanged: () => {} };
}

describe('ArbWsClient', () => {
  let fake: FakeSocket;
  let client: ArbWsClient;

  beforeEach(() => {
    vi.useFakeTimers();
    fake = new FakeSocket();
    client = new ArbWsClient({ socketFactory: () => fake });
    client.configure('http://127.0.0.1:7317', 'arb_token');
    client.connect();
    fake.fireOpen();
    fake.fireText({ type: 'hello_ok', protocol: 1, serverVersion: '1.8.0' });
  });

  afterEach(() => {
    vi.useRealTimers();
  });

  it('envoie hello puis sub après hello_ok', () => {
    expect(fake.controls('hello')).toHaveLength(1);
    expect(fake.controls('sub')[0]?.topics).toEqual(['sessions', 'worktrees', 'groups']);
  });

  it('corrèle attach→attached et délivre les OUTPUT au sink', async () => {
    const chunks: Uint8Array[] = [];
    const sink: AttachmentSink = { ...makeNoopSink(), data: (p) => chunks.push(p) };
    const p = client.attach({ sessionId: 'sess', mode: 'interactive', cols: 80, rows: 24, sink });
    fake.fireText({ type: 'attached', channel: 1, sessionId: 'sess', mode: 'interactive', controlling: true });
    const att = await p;
    expect(att.channel).toBe(1);

    fake.fireBinary(encodeBinaryFrame(BINARY_FRAME.OUTPUT, 1, new Uint8Array([104, 105]))); // "hi"
    expect(chunks).toHaveLength(1);
    expect([...(chunks[0] ?? [])]).toEqual([104, 105]);
  });

  it('ACK au seuil ACK_EVERY_BYTES', async () => {
    const p = client.attach({ sessionId: 'sess', mode: 'interactive', cols: 80, rows: 24, sink: makeNoopSink() });
    fake.fireText({ type: 'attached', channel: 7, sessionId: 'sess', mode: 'interactive', controlling: true });
    await p;
    fake.fireBinary(encodeBinaryFrame(BINARY_FRAME.OUTPUT, 7, new Uint8Array(FLOW.ACK_EVERY_BYTES)));
    const acks = fake.controls('ack');
    expect(acks).toHaveLength(1);
    expect(acks[0]).toMatchObject({ channel: 7, bytes: FLOW.ACK_EVERY_BYTES });
  });

  it('ACK traînant flushe un reliquat sous le seuil après 200ms', async () => {
    const p = client.attach({ sessionId: 'sess', mode: 'interactive', cols: 80, rows: 24, sink: makeNoopSink() });
    fake.fireText({ type: 'attached', channel: 2, sessionId: 'sess', mode: 'interactive', controlling: true });
    await p;
    fake.fireBinary(encodeBinaryFrame(BINARY_FRAME.OUTPUT, 2, new Uint8Array(1000)));
    expect(fake.controls('ack')).toHaveLength(0); // sous le seuil → pas d’ACK immédiat
    vi.advanceTimersByTime(200);
    expect(fake.controls('ack')).toEqual([{ type: 'ack', channel: 2, bytes: 1000 }]);
  });

  it('RESYNC réinitialise le terminal et les compteurs d’ACK', async () => {
    let resets = 0;
    const sink: AttachmentSink = { ...makeNoopSink(), reset: () => (resets += 1) };
    const p = client.attach({ sessionId: 'sess', mode: 'interactive', cols: 80, rows: 24, sink });
    fake.fireText({ type: 'attached', channel: 3, sessionId: 'sess', mode: 'interactive', controlling: true });
    await p;
    fake.fireBinary(encodeBinaryFrame(BINARY_FRAME.OUTPUT, 3, new Uint8Array(FLOW.ACK_EVERY_BYTES)));
    fake.fireBinary(encodeBinaryFrame(BINARY_FRAME.RESYNC, 3, new Uint8Array([65]))); // reset + replay "A"
    expect(resets).toBe(1);
    // après resync, le replay ne compte pas : un nouvel OUTPUT au seuil ré-ACK à partir de 0
    fake.fireBinary(encodeBinaryFrame(BINARY_FRAME.OUTPUT, 3, new Uint8Array(FLOW.ACK_EVERY_BYTES)));
    const acks = fake.controls('ack');
    expect(acks[acks.length - 1]).toMatchObject({ bytes: FLOW.ACK_EVERY_BYTES });
  });

  it('envoie stdin/resize sur le bon canal', async () => {
    const p = client.attach({ sessionId: 'sess', mode: 'interactive', cols: 80, rows: 24, sink: makeNoopSink() });
    fake.fireText({ type: 'attached', channel: 5, sessionId: 'sess', mode: 'interactive', controlling: true });
    const att = await p;
    att.sendStdin('ls\n');
    att.resize(120, 40);
    expect(fake.controls('stdin')[0]).toEqual({ type: 'stdin', channel: 5, data: 'ls\n' });
    expect(fake.controls('resize')[0]).toEqual({ type: 'resize', channel: 5, cols: 120, rows: 40 });
  });
});