train_utility/tests/frontend_smoke.js

const assert = require('node:assert/strict');
const path = require('node:path');

class FakeElement {
    constructor(id = '') {
        this.id = id;
        this.value = id === 'task' ? 'detect' : '';
        this.checked = false;
        this.disabled = false;
        this.style = {};
        this.children = [];
        this.listeners = {};
        this.className = '';
        this.textContent = '';
        this.scrollTop = 0;
        this.scrollHeight = 0;
        this.classList = {
            add() {},
            remove() {},
            contains() { return false; }
        };
    }

    addEventListener(name, handler) {
        this.listeners[name] = handler;
    }

    appendChild(child) {
        this.children.push(child);
        this.lastElementChild = child;
        return child;
    }

    getContext() {
        return {};
    }

    remove() {}

    set innerHTML(value) {
        this._innerHTML = value;
        this.children = [];
    }

    get innerHTML() {
        return this._innerHTML || '';
    }
}

const elements = new Map();
const element = id => {
    if (!elements.has(id)) elements.set(id, new FakeElement(id));
    return elements.get(id);
};

let domReady;
global.document = {
    addEventListener(name, handler) {
        if (name === 'DOMContentLoaded') domReady = handler;
    },
    querySelectorAll() {
        return [];
    },
    getElementById: element,
    createElement() {
        return new FakeElement();
    },
    head: new FakeElement('head'),
    body: new FakeElement('body')
};

const storage = new Map();
global.localStorage = {
    getItem(key) { return storage.has(key) ? storage.get(key) : null; },
    setItem(key, value) { storage.set(key, value); },
    removeItem(key) { storage.delete(key); }
};
global.confirm = () => true;
global.window = {
    location: {
        protocol: 'http:',
        host: '127.0.0.1:8000',
        hostname: '127.0.0.1'
    }
};

let nextTimerId = 1;
const scheduledTimers = [];
global.setTimeout = (handler, delay) => {
    const timer = {id: nextTimerId++, handler, delay, cancelled: false, fired: false};
    scheduledTimers.push(timer);
    return timer.id;
};
global.clearTimeout = id => {
    const timer = scheduledTimers.find(item => item.id === id);
    if (timer) timer.cancelled = true;
};

class FakeChart {
    static instances = [];

    constructor(_context, config) {
        this.data = config.data;
        this.options = config.options;
        FakeChart.instances.push(this);
    }

    destroy() {}
    update() {}
}
global.Chart = FakeChart;

class FakeWebSocket {
    static instances = [];

    constructor(url) {
        this.url = url;
        FakeWebSocket.instances.push(this);
    }
}
global.WebSocket = FakeWebSocket;

const response = (ok, data) => ({
    ok,
    async json() { return data; }
});
const fetchCalls = [];
let holdTrainingStart = false;
let holdExportStart = false;
global.fetch = async (url, options = {}) => {
    fetchCalls.push({url, options});
    if (url === '/api/train/start' && holdTrainingStart) {
        return new Promise(() => {});
    }
    if (url === '/api/export/start' && holdExportStart) {
        return new Promise(() => {});
    }
    if (url === '/api/datasets' || url === '/api/models' || url === '/api/sessions') {
        if (url === '/api/models') {
            return response(true, [{
                name: 'detect/experiment/best.pt',
                path: '/workspace/runs/detect/experiment/weights/best.pt',
                source: 'runs'
            }]);
        }
        return response(true, []);
    }
    if (url === '/api/sessions/last_run') return response(false, {});
    if (url === '/api/config/defaults') {
        return response(true, {
            dataset: 'coco8.yaml',
            model: 'yolo11n.pt',
            task: 'detect',
            workers: 8,
            patience: 100,
            augmentation: {enabled: true, close_mosaic: 10},
            mlflow: {enabled: false, tracking_uri: 'http://mlflow:5000'},
            split: {enabled: false}
        });
    }
    return response(false, {});
};

require(path.resolve(__dirname, '../src/yolo_webui/static/app.js'));

async function flushPromises() {
    await new Promise(resolve => setImmediate(resolve));
    await new Promise(resolve => setImmediate(resolve));
}

function descendants(node) {
    return node.children.flatMap(child => [child, ...descendants(child)]);
}

function runTimer(timer) {
    timer.fired = true;
    timer.handler();
}

(async () => {
    assert.equal(typeof domReady, 'function');
    domReady();
    await flushPromises();

    assert.equal(element('mlflow-header-link').href, 'http://127.0.0.1:5000/');

    const localModelPath = '/workspace/runs/detect/experiment/weights/best.pt';
    const localModelOption = descendants(element('model-select'))
        .find(item => item.textContent === 'detect/experiment/best.pt');
    assert.ok(localModelOption);
    assert.equal(localModelOption.value, localModelPath);

    element('model-select').value = localModelOption.value;
    element('model-select').listeners.change();
    element('workers').value = '0';
    element('patience').value = '0';
    element('close-mosaic').value = '0';
    element('auto-augment').value = 'none';
    element('config-form').listeners.input();
    const savedConfig = JSON.parse(storage.get('draft_config'));
    assert.equal(savedConfig.workers, 0);
    assert.equal(savedConfig.patience, 0);
    assert.equal(savedConfig.augmentation.close_mosaic, 0);
    assert.equal(savedConfig.augmentation.auto_augment, 'none');
    assert.equal(savedConfig.model, localModelPath);

    element('model-select').value = '__custom__';
    element('model-select').listeners.change();
    element('model').value = '/workspace/models/custom-segment.pt';
    element('task').value = 'segment';
    element('task').listeners.change();
    assert.equal(element('model-select').value, '__custom__');
    assert.equal(element('model').value, '/workspace/models/custom-segment.pt');

    assert.equal(FakeWebSocket.instances.length, 2);
    const socket = FakeWebSocket.instances[0];
    const exportSocket = FakeWebSocket.instances[1];
    const originalConsoleError = console.error;
    console.error = () => {};
    try {
        assert.doesNotThrow(() => socket.onmessage({data: '{broken'}));
        assert.doesNotThrow(() => exportSocket.onmessage({data: 'null'}));
    } finally {
        console.error = originalConsoleError;
    }

    socket.onmessage({
        data: JSON.stringify({
            type: 'init',
            status: 'idle',
            logs: [],
            metrics: [
                {epoch: 1, mAP50: 0.5},
                {epoch: 2, loss: 0.2}
            ]
        })
    });

    const chart = FakeChart.instances.at(-1);
    assert.deepEqual(chart.data.labels, [1, 2]);
    assert.deepEqual(chart.data.datasets.map(item => item.label), ['mAP50', 'loss']);
    assert.deepEqual(chart.data.datasets[0].data, [0.5, null]);
    assert.deepEqual(chart.data.datasets[1].data, [null, 0.2]);

    socket.onclose();
    socket.onclose();
    let activeReconnectTimers = scheduledTimers.filter(
        timer => timer.delay === 5000 && !timer.cancelled && !timer.fired
    );
    assert.equal(activeReconnectTimers.length, 1);
    runTimer(activeReconnectTimers[0]);
    assert.equal(FakeWebSocket.instances.length, 3);
    assert.equal(FakeWebSocket.instances[2].url, 'ws://127.0.0.1:8000/api/ws');

    exportSocket.onclose();
    exportSocket.onclose();
    activeReconnectTimers = scheduledTimers.filter(
        timer => timer.delay === 5000 && !timer.cancelled && !timer.fired
    );
    assert.equal(activeReconnectTimers.length, 1);
    runTimer(activeReconnectTimers[0]);
    assert.equal(FakeWebSocket.instances.length, 4);
    assert.equal(FakeWebSocket.instances[3].url, 'ws://127.0.0.1:8000/api/export/ws');

    element('session-name').value = 'profile with spaces';
    await element('session-save-btn').listeners.click();
    assert.equal(
        fetchCalls.filter(call => call.options.method === 'POST' && call.url.startsWith('/api/sessions/')).length,
        0
    );

    element('export-model').value = localModelPath;
    element('export-format').value = 'pb';
    element('export-imgsz').value = '31';
    element('export-batch').value = '1';
    element('export-workspace').value = '4';
    await element('export-start-btn').listeners.click();
    assert.equal(fetchCalls.filter(call => call.url === '/api/export/start').length, 0);

    element('export-imgsz').value = '640';
    element('export-workspace').value = '1.5';
    element('export-half').checked = true;
    element('export-int8').checked = true;
    await element('export-start-btn').listeners.click();
    assert.equal(fetchCalls.filter(call => call.url === '/api/export/start').length, 0);

    element('export-half').checked = false;
    element('export-int8').checked = false;
    holdExportStart = true;
    element('export-start-btn').listeners.click();
    element('export-start-btn').listeners.click();
    const exportStartCalls = fetchCalls.filter(call => call.url === '/api/export/start');
    assert.equal(exportStartCalls.length, 1);
    assert.equal(element('export-start-btn').disabled, true);
    assert.deepEqual(
        JSON.parse(exportStartCalls[0].options.body),
        {
            model: localModelPath,
            format: 'pb',
            imgsz: 640,
            half: false,
            int8: false,
            dynamic: false,
            simplify: false,
            batch: 1,
            workspace: 1.5
        }
    );

    holdTrainingStart = true;
    element('start-btn').listeners.click();
    element('start-btn').listeners.click();
    assert.equal(fetchCalls.filter(call => call.url === '/api/train/start').length, 1);
    assert.equal(element('start-btn').disabled, true);
})().catch(error => {
    console.error(error);
    process.exitCode = 1;
});