Mujoco_WASM/web_platform/src/training/LocalTrainingPanel.test.tsx

import { fireEvent, render, screen, waitFor } from '@testing-library/react';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { LocalTrainingPanel } from './LocalTrainingPanel';

beforeEach(() => {
  vi.restoreAllMocks();
  localStorage.clear();
  sessionStorage.clear();
  vi.unstubAllGlobals();
});

describe('LocalTrainingPanel', () => {
  it('连接本地服务并从图形界面发起训练请求', async () => {
    const health = {
      version: '0.1.0',
      ready: true,
      trainerRoot: '/opt/unitree_rl_mjlab',
      python: '/env/bin/python',
      tasks: ['Unitree-Go2-Flat'],
    };
    const job = {
      id: 'a'.repeat(32),
      state: 'queued',
      taskId: 'Unitree-Go2-Flat',
      createdAt: '2025-01-01T00:00:00Z',
      iteration: 0,
      maxIterations: 2000,
      progress: 0,
      message: '等待启动',
      logs: [],
      artifactReady: false,
    };
    const fetchMock = vi
      .fn()
      .mockResolvedValueOnce(
        new Response(JSON.stringify(health), {
          status: 200,
          headers: { 'Content-Type': 'application/json' },
        }),
      )
      .mockResolvedValueOnce(
        new Response(JSON.stringify({ presets: [] }), {
          status: 200,
          headers: { 'Content-Type': 'application/json' },
        }),
      )
      .mockResolvedValueOnce(
        new Response(JSON.stringify(job), {
          status: 202,
          headers: { 'Content-Type': 'application/json' },
        }),
      )
      .mockResolvedValueOnce(
        new Response(JSON.stringify(health), {
          status: 200,
          headers: { 'Content-Type': 'application/json' },
        }),
      )
      .mockResolvedValueOnce(
        new Response(JSON.stringify({ presets: [] }), {
          status: 200,
          headers: { 'Content-Type': 'application/json' },
        }),
      )
      .mockResolvedValueOnce(
        new Response(JSON.stringify({ error: '训练任务不存在或服务已重启' }), {
          status: 404,
          headers: { 'Content-Type': 'application/json' },
        }),
      );
    vi.stubGlobal('fetch', fetchMock);
    render(<LocalTrainingPanel onPolicyReady={vi.fn()} />);
    fireEvent.change(screen.getByLabelText('训练服务访问令牌'), {
      target: { value: 'secret-token' },
    });
    fireEvent.click(screen.getByRole('button', { name: '连接' }));
    expect(await screen.findByText('/opt/unitree_rl_mjlab')).toBeInTheDocument();
    const open = vi.spyOn(window, 'open').mockImplementation(() => null);
    fireEvent.click(screen.getByRole('button', { name: '打开自调参 Agent 工作台' }));
    expect(open).toHaveBeenCalledWith(expect.any(URL), '_blank');
    expect(String(open.mock.calls[0][0])).toContain('tuning.html');
    expect(String(open.mock.calls[0][0])).not.toContain('secret-token');
    fireEvent.change(screen.getByLabelText('并行环境'), { target: { value: '32' } });
    fireEvent.click(screen.getByRole('button', { name: '发起本地训练' }));
    await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(3));
    const request = fetchMock.mock.calls[2][1] as RequestInit;
    expect(JSON.parse(String(request.body))).toMatchObject({
      taskId: 'Unitree-Go2-Flat',
      numEnvs: 32,
      device: 'gpu',
      gpuIds: [0],
      wandbMode: 'offline',
    });
    expect(new Headers(request.headers).get('Authorization')).toBe('Bearer secret-token');
    expect(await screen.findByText('排队中')).toBeInTheDocument();

    const tokenInput = screen.getByLabelText('训练服务访问令牌');
    expect(tokenInput).toBeEnabled();
    fireEvent.change(tokenInput, { target: { value: 'new-secret-token' } });
    fireEvent.click(screen.getByRole('button', { name: '连接' }));
    await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(6));
    expect(await screen.findByRole('button', { name: '发起本地训练' })).toBeInTheDocument();
    expect(sessionStorage.getItem('mujoco-local-training-token')).toBe('new-secret-token');
  });

  it('接收调参窗口传回的策略文件，无需主工作台重复持有令牌', async () => {
    const onPolicyReady = vi.fn();
    const reply = vi.spyOn(window, 'postMessage').mockImplementation(() => undefined);
    render(<LocalTrainingPanel onPolicyReady={onPolicyReady} />);
    const policy = new File([new Uint8Array([1, 2, 3])], 'best-policy.onnx', {
      type: 'application/octet-stream',
    });

    window.dispatchEvent(
      new MessageEvent('message', {
        origin: window.location.origin,
        source: window,
        data: {
          type: 'mujoco-tuning-import-policy',
          sessionId: 'a'.repeat(32),
          policy,
        },
      }),
    );

    await waitFor(() => expect(onPolicyReady).toHaveBeenCalledWith(policy));
    expect(reply).toHaveBeenCalledWith(
      expect.objectContaining({
        type: 'mujoco-tuning-import-policy-result',
        ok: true,
      }),
      window.location.origin,
    );
  });

  it('旧调参消息缺少主工作台令牌时显示错误而不是抛出未处理异常', async () => {
    const reply = vi.spyOn(window, 'postMessage').mockImplementation(() => undefined);
    render(<LocalTrainingPanel onPolicyReady={vi.fn()} />);

    window.dispatchEvent(
      new MessageEvent('message', {
        origin: window.location.origin,
        source: window,
        data: {
          type: 'mujoco-tuning-import-policy',
          sessionId: 'a'.repeat(32),
        },
      }),
    );

    expect(await screen.findByRole('alert')).toHaveTextContent('请输入训练服务访问令牌');
    expect(reply).toHaveBeenCalledWith(
      expect.objectContaining({
        type: 'mujoco-tuning-import-policy-result',
        ok: false,
        error: '请输入训练服务访问令牌',
      }),
      window.location.origin,
    );
  });
});