Shade/packages/shade-files/tests/integration/custom-op.test.ts

import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
import { z } from 'zod';
import { setupFileRig, type FileTestRig } from './helpers/rig.js';
import { CustomOpRejectedError, NotImplementedError } from '../../src/index.js';

// Module augmentation: register a typed custom op for the test.
declare module '../../src/index.js' {
  interface CustomOpsMap {
    'test.echo': { args: { message: string }; response: { echoed: string } };
    'test.add': { args: { a: number; b: number }; response: { sum: number } };
  }
}

describe('Custom ops — registry + Zod validation + typed I/O', () => {
  let rig: FileTestRig;
  const callLog: string[] = [];

  beforeAll(async () => {
    rig = await setupFileRig({
      custom: {
        'test.echo': {
          args: z.object({ message: z.string().min(1).max(64) }),
          response: z.object({ echoed: z.string() }),
          handler: async (args, ctx) => {
            callLog.push(`echo:${ctx.sender}:${args.message}`);
            return { echoed: args.message.toUpperCase() };
          },
        },
        'test.add': {
          args: z.object({ a: z.number(), b: z.number() }),
          response: z.object({ sum: z.number() }),
          handler: async (args) => ({ sum: args.a + args.b }),
        },
        'test.bad-response': {
          args: z.object({}),
          response: z.object({ x: z.number() }),
          // Returns wrong shape on purpose.
          handler: async () => ({ y: 'not-a-number' }) as unknown as { x: number },
        },
      },
    });
  });

  afterAll(async () => {
    await rig.teardown();
  });

  test('typed echo round-trips through registered Zod schemas', async () => {
    const result = await rig.fs.custom('test.echo', { message: 'hello' });
    expect(result.echoed).toBe('HELLO');
    expect(callLog).toContain('echo:alice:hello');
  });

  test('typed add', async () => {
    const result = await rig.fs.custom('test.add', { a: 3, b: 4 });
    expect(result.sum).toBe(7);
  });

  test('invalid args (Zod-rejected payload) → InvalidArgsError', async () => {
    await expect(
      // message: '' violates min(1) — TypeScript still allows it since string
      rig.fs.custom('test.echo', { message: '' }),
    ).rejects.toThrow();
  });

  test('unknown custom op name → NotImplementedError', async () => {
    await expect(
      rig.fs.custom('test.unknown' as never, {} as never),
    ).rejects.toBeInstanceOf(NotImplementedError);
  });

  test('handler returns wrong shape → CustomOpRejectedError', async () => {
    await expect(
      rig.fs.custom('test.bad-response' as never, {} as never),
    ).rejects.toBeInstanceOf(CustomOpRejectedError);
  });
});