wisp/src/lib/transcription/pipeline.test.ts

import { describe, it, expect } from 'vitest';
import { transcribe } from './pipeline';
import type { TranscribeProgress } from './pipeline';
import type { EngineCapabilities, TranscriptionEngine } from './engine';
import type {
  ModelId,
  PcmAudio,
  Segment,
  TranscribeOptions,
} from '../types';
import { WHISPER_SAMPLE_RATE } from '../types';
import { planChunks } from './chunking';

/**
 * A fully in-memory fake engine. It records calls and returns exactly one
 * segment per chunk so the test can assert how the pipeline offsets/stitches
 * chunk-local times into absolute time.
 */
class FakeEngine implements TranscriptionEngine {
  readonly platform = 'web' as const;

  /** Toggles whether loadModel is exercised. */
  loaded: boolean;
  loadModelCalls = 0;
  transcribeCalls = 0;
  lastLoadProgress: number[] = [];

  constructor(opts: { loaded: boolean } = { loaded: true }) {
    this.loaded = opts.loaded;
  }

  async capabilities(): Promise<EngineCapabilities> {
    return {
      backend: 'wasm',
      supportsLiveMic: false,
      maxRecommendedModel: 'tiny.en',
    };
  }

  async loadModel(
    _modelId: ModelId,
    onProgress?: (p: number) => void,
  ): Promise<void> {
    this.loadModelCalls++;
    // Report a couple of fractional progress ticks then complete.
    onProgress?.(0.5);
    onProgress?.(1);
    this.lastLoadProgress.push(1);
    this.loaded = true;
  }

  isModelLoaded(_modelId: ModelId): boolean {
    return this.loaded;
  }

  async transcribeChunk(
    audio: PcmAudio,
    _opts: TranscribeOptions,
  ): Promise<Segment[]> {
    // One chunk-local segment per chunk, labelled by call index so we can
    // verify ordering and absolute-time offsetting after stitching. We place it
    // squarely in the chunk's interior (mid-window) so the stitcher's
    // mid-overlap cut keeps exactly one copy per chunk rather than dropping a
    // boundary-hugging segment. Times are chunk-LOCAL (the engine contract).
    const idx = this.transcribeCalls++;
    const lenSec = audio.samples.length / audio.sampleRate;
    const mid = lenSec / 2;
    return [{ start: mid, end: mid + 0.5, text: `chunk${idx}` }];
  }
}

/** Build a silent PcmAudio of `seconds` length at 16 kHz. */
function makeAudio(seconds: number): PcmAudio {
  return {
    sampleRate: WHISPER_SAMPLE_RATE,
    samples: new Float32Array(Math.round(seconds * WHISPER_SAMPLE_RATE)),
  };
}

const OPTIONS: TranscribeOptions = { modelId: 'tiny.en' };

describe('transcribe pipeline', () => {
  it('returns stitched, absolute-time segments (one per chunk)', async () => {
    // 60s with the default 28s window / 2s overlap -> multiple chunks.
    const audio = makeAudio(60);
    const engine = new FakeEngine({ loaded: true });
    const plan = planChunks(audio.samples.length, {
      windowSec: 28,
      overlapSec: 2,
      sampleRate: WHISPER_SAMPLE_RATE,
    });

    const segments = await transcribe({ audio, options: OPTIONS, engine });

    // One segment per chunk survives stitching (each is well inside its window).
    expect(engine.transcribeCalls).toBe(plan.length);
    expect(segments).toHaveLength(plan.length);

    // Each segment was offset by its chunk's media start time. The fake emits a
    // chunk-local segment at the window midpoint; after offsetting, segment i
    // should start at chunk.startSec + (window/2), i.e. strictly inside chunk i.
    for (let i = 0; i < segments.length; i++) {
      const seg = segments[i]!;
      const chunk = plan[i]!;
      const localMid = (chunk.endSec - chunk.startSec) / 2;
      expect(seg.text).toBe(`chunk${i}`);
      expect(seg.start).toBeCloseTo(chunk.startSec + localMid, 5);
    }
    // Output must be sorted ascending by start.
    for (let i = 1; i < segments.length; i++) {
      expect(segments[i]!.start).toBeGreaterThanOrEqual(segments[i - 1]!.start);
    }
  });

  it('reports increasing transcribing progress ending at 1', async () => {
    const audio = makeAudio(60);
    const engine = new FakeEngine({ loaded: true });
    const events: TranscribeProgress[] = [];

    await transcribe({
      audio,
      options: OPTIONS,
      engine,
      onProgress: (p) => events.push(p),
    });

    const transcribing = events.filter((e) => e.stage === 'transcribing');
    expect(transcribing.length).toBeGreaterThan(0);

    // The first transcribing event is a kickoff fired BEFORE the first chunk so
    // the UI leaves the "loading" stage immediately: progress 0, no partial yet.
    expect(transcribing[0]!.progress).toBe(0);
    expect(transcribing[0]!.partial).toEqual([]);

    // Every event carries the total chunk count; progress stays within [0, 1].
    for (const e of transcribing) {
      expect(e.chunkCount).toBeGreaterThan(0);
      expect(e.progress).toBeGreaterThanOrEqual(0);
      expect(e.progress).toBeLessThanOrEqual(1);
    }
    // After the kickoff, per-chunk progress strictly increases and ends at 1.
    for (let i = 2; i < transcribing.length; i++) {
      expect(transcribing[i]!.progress).toBeGreaterThan(transcribing[i - 1]!.progress);
    }
    expect(transcribing[transcribing.length - 1]!.progress).toBe(1);

    // The growing partial should reach the final segment count on the last tick.
    const lastPartial = transcribing[transcribing.length - 1]!.partial;
    expect(lastPartial.length).toBeGreaterThan(0);
  });

  it('calls loadModel (with load progress) when the model is not loaded', async () => {
    const audio = makeAudio(10); // single chunk is fine here
    const engine = new FakeEngine({ loaded: false });
    const loadingEvents: TranscribeProgress[] = [];

    await transcribe({
      audio,
      options: OPTIONS,
      engine,
      onProgress: (p) => {
        if (p.stage === 'loading') loadingEvents.push(p);
      },
    });

    expect(engine.loadModelCalls).toBe(1);
    // Loading events have an empty partial and a fraction in [0, 1].
    expect(loadingEvents.length).toBeGreaterThan(0);
    for (const e of loadingEvents) {
      expect(e.partial).toEqual([]);
      expect(e.progress).toBeGreaterThanOrEqual(0);
      expect(e.progress).toBeLessThanOrEqual(1);
    }
    // Last loading tick reaches 1.
    expect(loadingEvents[loadingEvents.length - 1]!.progress).toBe(1);
  });

  it('does NOT call loadModel when the model is already loaded', async () => {
    const audio = makeAudio(10);
    const engine = new FakeEngine({ loaded: true });
    await transcribe({ audio, options: OPTIONS, engine });
    expect(engine.loadModelCalls).toBe(0);
  });

  it('throws AbortError when the signal is already aborted', async () => {
    const audio = makeAudio(60);
    const engine = new FakeEngine({ loaded: true });
    const controller = new AbortController();
    controller.abort();

    await expect(
      transcribe({
        audio,
        options: OPTIONS,
        engine,
        signal: controller.signal,
      }),
    ).rejects.toMatchObject({ name: 'AbortError' });

    // Nothing was transcribed because we aborted before the first chunk.
    expect(engine.transcribeCalls).toBe(0);
  });
});