SnowLuma/packages/core/tests/webui-debug-tools.test.ts

import { afterAll, describe, expect, it } from 'vitest';
import { readdir, readFile, rm } from 'fs/promises';
import path from 'path';
import {
    buildStreamInvokeResponse,
    DEBUG_UPLOAD_DIR,
    safeUploadName,
    streamUploadToDisk,
    uploadDestPath,
} from '../src/webui/debug-tools';

const written: string[] = [];
afterAll(async () => {
    await Promise.all(written.map((p) => rm(p, { force: true }).catch(() => {})));
});

function streamOf(...chunks: Uint8Array[]): ReadableStream<Uint8Array> {
    let i = 0;
    return new ReadableStream<Uint8Array>({
        pull(controller) {
            if (i < chunks.length) {
                controller.enqueue(chunks[i++]);
            } else {
                controller.close();
            }
        },
    });
}

async function readSse(res: Response): Promise<string[]> {
    const text = await res.text();
    return text
        .split('\n\n')
        .filter((s) => s.startsWith('data: '))
        .map((s) => s.slice('data: '.length));
}

describe('safeUploadName', () => {
    it('strips path traversal and separators to a bare basename', () => {
        expect(safeUploadName('../../etc/passwd')).toBe('passwd');
        expect(safeUploadName('/abs/secret.key')).toBe('secret.key');
        expect(safeUploadName('a\\b\\c.png')).toBe('c.png');
        expect(safeUploadName('..')).toBe('file'); // leading dots stripped → empty → fallback
        expect(safeUploadName('...hidden')).toBe('hidden');
    });
    it('collapses exotic chars and falls back for empties', () => {
        expect(safeUploadName('my pic!@#.jpg')).toBe('my_pic_.jpg');
        expect(safeUploadName('')).toBe('file');
        expect(safeUploadName(undefined)).toBe('file');
    });
    it('keeps the extension and caps length', () => {
        const long = 'x'.repeat(500) + '.mp4';
        const out = safeUploadName(long);
        expect(out.length).toBeLessThanOrEqual(120);
        expect(out.startsWith('x')).toBe(true);
    });
});

describe('uploadDestPath', () => {
    it('uses a neutral temporary directory name', () => {
        expect(path.basename(DEBUG_UPLOAD_DIR)).toBe('webui-upload');
    });

    it('produces a unique path inside the upload dir', () => {
        const a = uploadDestPath('pic.png');
        const b = uploadDestPath('pic.png');
        expect(a).not.toBe(b); // random id prefix
        expect(path.dirname(a)).toBe(DEBUG_UPLOAD_DIR);
        expect(a.endsWith('__pic.png')).toBe(true);
        const rel = path.relative(DEBUG_UPLOAD_DIR, a);
        expect(rel.startsWith('..')).toBe(false);
        expect(path.isAbsolute(rel)).toBe(false);
    });
});

describe('streamUploadToDisk', () => {
    it('streams bytes to disk and reports the size', async () => {
        const data = new TextEncoder().encode('hello debug upload');
        const res = await streamUploadToDisk(streamOf(data.slice(0, 5), data.slice(5)), 'note.txt');
        written.push(res.path);
        expect(res.size).toBe(data.byteLength);
        expect(path.dirname(res.path)).toBe(DEBUG_UPLOAD_DIR);
        expect(await readFile(res.path, 'utf8')).toBe('hello debug upload');
    });

    it('rejects and removes the partial file when the cap is exceeded', async () => {
        const big = new Uint8Array(64);
        await expect(
            streamUploadToDisk(streamOf(big, big), 'big.bin', { maxBytes: 100 }),
        ).rejects.toThrow(/大小上限/);
        // the partial must not survive — no leftover `__big.bin` in the dir
        const names = await readdir(DEBUG_UPLOAD_DIR).catch(() => [] as string[]);
        expect(names.filter((n) => n.endsWith('__big.bin'))).toEqual([]);
    });

    it('rejects an empty body', async () => {
        await expect(streamUploadToDisk(null, 'x')).rejects.toThrow(/empty/);
        await expect(streamUploadToDisk(streamOf(), 'x')).rejects.toThrow(/empty/);
    });

    it('propagates a mid-transfer source error and removes the partial', async () => {
        const failing = new ReadableStream<Uint8Array>({
            start(controller) {
                controller.enqueue(new TextEncoder().encode('partial-bytes'));
                controller.error(new Error('source aborted'));
            },
        });
        await expect(streamUploadToDisk(failing, 'oops.bin')).rejects.toThrow(/source aborted/);
        const names = await readdir(DEBUG_UPLOAD_DIR).catch(() => [] as string[]);
        expect(names.filter((n) => n.endsWith('__oops.bin'))).toEqual([]);
    });
});

describe('buildStreamInvokeResponse', () => {
    const noDelay = async () => {};

    it('cancels delivery while the action is still pending', async () => {
        let finish!: () => void;
        const pending = new Promise<void>((resolve) => {
            finish = resolve;
        });
        let isAlive!: () => boolean;
        const response = buildStreamInvokeResponse(
            async (_request, _emit, alive) => {
                isAlive = alive;
                await pending;
            },
            '{}',
            new AbortController().signal,
        );
        const canceled = response.body!.cancel();

        try {
            expect(isAlive()).toBe(false);
        } finally {
            finish();
            await canceled;
        }
    });

    it('does not invoke an action for an already aborted request', async () => {
        const controller = new AbortController();
        controller.abort();
        let invoked = false;
        const response = buildStreamInvokeResponse(
            async () => {
                invoked = true;
            },
            '{}',
            controller.signal,
        );
        await readSse(response);
        expect(invoked).toBe(false);
    });

    it('relays each emitted frame as one SSE event, in order', async () => {
        const ac = new AbortController();

        const driver = async (_rq: string, emit: (j: string) => Promise<void>) => {
            await emit(JSON.stringify({ data: { type: 'stream', i: 1 } }));
            await emit(JSON.stringify({ data: { type: 'stream', i: 2 } }));
            await emit(JSON.stringify({ status: 'ok', data: { type: 'response' } }));
        };

        const res = buildStreamInvokeResponse(driver, '{"action":"x"}', ac.signal, noDelay);
        const frames = (await readSse(res)).map((s) => JSON.parse(s));
        expect(frames).toHaveLength(3);
        expect(frames[0].data.i).toBe(1);
        expect(frames[2].data.type).toBe('response');
    });

    it('turns a driver error into a terminal failed frame', async () => {
        const ac = new AbortController();

        const driver = async (_rq: string, emit: (j: string) => Promise<void>) => {
            await emit(JSON.stringify({ data: { type: 'stream' } }));
            throw new Error('disk exploded');
        };

        const res = buildStreamInvokeResponse(driver, '{}', ac.signal, noDelay);
        const frames = (await readSse(res)).map((s) => JSON.parse(s));
        expect(frames).toHaveLength(2);
        expect(frames[1].status).toBe('failed');
        expect(frames[1].message).toMatch(/disk exploded/);
    });

    it('passes the raw request through to the driver', async () => {
        const ac = new AbortController();
        let seen = '';

        const driver = async (rq: string, emit: (j: string) => Promise<void>) => {
            seen = rq;
            await emit('{"data":{"type":"response"}}');
        };

        const res = buildStreamInvokeResponse(
            driver,
            '{"action":"test_download_stream"}',
            ac.signal,
            noDelay,
        );
        await readSse(res);
        expect(seen).toBe('{"action":"test_download_stream"}');
    });

    it('exposes an SSE content-type', () => {
        const ac = new AbortController();
        const res = buildStreamInvokeResponse(async () => {}, '{}', ac.signal, noDelay);
        expect(res.headers.get('Content-Type')).toBe('text/event-stream');
    });
});