Portfolio/tests/unit/shader-warmup.test.ts

import { describe, expect, it } from 'vitest';
import { warmUpShaders, type WarmupRenderer, type WarmupScheduler } from '@/scene/world/shader-warmup';

interface Node {
  name: string;
  visible: boolean;
  children: Node[];
  traverse(callback: (object: Node) => void): void;
}

const node = (name: string, visible: boolean, children: Node[] = []): Node => ({
  name,
  visible,
  children,
  traverse(callback) {
    callback(this);
    for (const child of this.children) child.traverse(callback);
  },
});

const manualScheduler = () => {
  let clock = 0;
  const timers = new Map<number, () => void>();
  let next = 1;
  const scheduler: WarmupScheduler = {
    now: () => clock,
    setTimeout(callback, delay) {
      const id = next++;
      timers.set(id, () => {
        clock += delay;
        callback();
      });
      return id;
    },
    clearTimeout(handle) {
      timers.delete(handle as number);
    },
  };
  return { scheduler, fire: () => [...timers.values()].forEach((run) => run()), pending: () => timers.size };
};

const fakeRenderer = (compile: () => Promise<unknown>, parallel = true) => {
  const seen: { target: unknown; visible: string[] }[] = [];
  const state = { target: 'screen-previous' as unknown };
  let scene: Node | null = null;
  const renderer: WarmupRenderer = {
    getRenderTarget: () => state.target,
    setRenderTarget: (target) => {
      state.target = target;
    },
    compileAsync: (root) => {
      scene = root as Node;
      const visible: string[] = [];
      scene.traverse((object) => {
        if (object.visible) visible.push(object.name);
      });
      seen.push({ target: state.target, visible });
      return compile();
    },
    extensions: { has: (name) => parallel && name === 'KHR_parallel_shader_compile' },
  };
  return { renderer, seen, state };
};

describe('shader warm-up', () => {
  it('compiles against the scene target with every subtree visible, then restores both', async () => {
    const lamp = node('hidden-light', false);
    const scene = node('root', true, [node('car', true), node('cone', false, [lamp])]);
    const { renderer, seen, state } = fakeRenderer(() => Promise.resolve());
    const { scheduler } = manualScheduler();

    const result = await warmUpShaders(renderer, scene, {}, 'composer-target', 2_500, scheduler);

    expect(result).toMatchObject({ outcome: 'ready', parallel: true });
    expect(seen).toEqual([
      { target: 'composer-target', visible: ['root', 'car', 'cone', 'hidden-light'] },
    ]);
    expect(state.target).toBe('screen-previous');
    expect(scene.children[1]!.visible).toBe(false);
    expect(lamp.visible).toBe(false);
  });

  it('gives up after the timeout so a slow driver never holds the build', async () => {
    const { renderer } = fakeRenderer(() => new Promise(() => undefined));
    const timers = manualScheduler();
    const pending = warmUpShaders(renderer, node('root', true), {}, null, 2_500, timers.scheduler);
    timers.fire();
    await expect(pending).resolves.toMatchObject({ outcome: 'timeout', waited: 2_500 });
  });

  it('clears its timer once the compile resolves', async () => {
    const { renderer } = fakeRenderer(() => Promise.resolve(), false);
    const timers = manualScheduler();
    const result = await warmUpShaders(renderer, node('root', true), {}, null, 2_500, timers.scheduler);
    expect(result).toMatchObject({ outcome: 'ready', parallel: false });
    expect(timers.pending()).toBe(0);
  });

  it('reports a throwing compile as failed and still restores the renderer and scene', async () => {
    const hidden = node('hidden', false);
    const { renderer, state } = fakeRenderer(() => {
      throw new Error('link failed');
    });
    const result = await warmUpShaders(renderer, node('root', true, [hidden]), {}, 'target', 2_500, manualScheduler().scheduler);
    expect(result.outcome).toBe('failed');
    expect(state.target).toBe('screen-previous');
    expect(hidden.visible).toBe(false);
  });

  it('reports a rejected compile as failed', async () => {
    const { renderer } = fakeRenderer(() => Promise.reject(new Error('lost context')));
    const result = await warmUpShaders(renderer, node('root', true), {}, null, 2_500, manualScheduler().scheduler);
    expect(result.outcome).toBe('failed');
  });
});