Pre_Seg_Server/src/components/VideoWorkspace.test.tsx

import { act, fireEvent, render, screen, waitFor } from '@testing-library/react';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { resetStore } from '../test/storeTestUtils';
import { useStore } from '../store/useStore';
import { VideoWorkspace } from './VideoWorkspace';

const apiMock = vi.hoisted(() => ({
  getProjectFrames: vi.fn(),
  parseMedia: vi.fn(),
  propagateMasks: vi.fn(),
  getTask: vi.fn(),
  getTemplates: vi.fn(),
  getProjectAnnotations: vi.fn(),
  saveAnnotation: vi.fn(),
  updateAnnotation: vi.fn(),
  deleteAnnotation: vi.fn(),
  exportCoco: vi.fn(),
  exportMasks: vi.fn(),
  importGtMask: vi.fn(),
  annotationToMask: vi.fn(),
  buildAnnotationPayload: vi.fn(),
  getAiModelStatus: vi.fn(),
}));

vi.mock('../lib/api', () => ({
  getProjectFrames: apiMock.getProjectFrames,
  parseMedia: apiMock.parseMedia,
  propagateMasks: apiMock.propagateMasks,
  getTask: apiMock.getTask,
  getTemplates: apiMock.getTemplates,
  getProjectAnnotations: apiMock.getProjectAnnotations,
  saveAnnotation: apiMock.saveAnnotation,
  updateAnnotation: apiMock.updateAnnotation,
  deleteAnnotation: apiMock.deleteAnnotation,
  exportCoco: apiMock.exportCoco,
  exportMasks: apiMock.exportMasks,
  importGtMask: apiMock.importGtMask,
  annotationToMask: apiMock.annotationToMask,
  buildAnnotationPayload: apiMock.buildAnnotationPayload,
  getAiModelStatus: apiMock.getAiModelStatus,
}));

describe('VideoWorkspace', () => {
  beforeEach(() => {
    resetStore();
    vi.clearAllMocks();
    useStore.setState({ currentProject: { id: '1', name: 'Demo', status: 'ready', video_path: 'uploads/demo.mp4' } });
    apiMock.getTemplates.mockResolvedValue([]);
    apiMock.getProjectAnnotations.mockResolvedValue([]);
    apiMock.annotationToMask.mockReturnValue(null);
    apiMock.getTask.mockResolvedValue({ id: 1, status: 'success', progress: 100, message: '解析完成' });
    apiMock.propagateMasks.mockResolvedValue({
      model: 'sam2',
      direction: 'forward',
      source_frame_id: 10,
      processed_frame_count: 3,
      created_annotation_count: 2,
      annotations: [],
    });
    apiMock.getAiModelStatus.mockResolvedValue({
      selected_model: 'sam2',
      gpu: { available: false, device: 'cpu', name: null, torch_available: true },
      models: [
        { id: 'sam2', label: 'SAM 2', available: true, loaded: false, device: 'cpu', supports: [], message: 'ready', package_available: true, checkpoint_exists: true, python_ok: true, torch_ok: true, cuda_required: false },
        { id: 'sam3', label: 'SAM 3', available: false, loaded: false, device: 'unavailable', supports: [], message: 'missing', package_available: false, checkpoint_exists: false, python_ok: false, torch_ok: true, cuda_required: true },
      ],
    });
  });

  it('loads project frames into the workspace store', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);

    render(<VideoWorkspace />);

    await waitFor(() => expect(useStore.getState().frames).toEqual([
      { id: '10', projectId: '1', index: 0, url: '/frame.jpg', width: 640, height: 360 },
    ]));
    expect(screen.getByText('Demo')).toBeInTheDocument();
    expect(apiMock.parseMedia).not.toHaveBeenCalled();
    expect(apiMock.getProjectAnnotations).toHaveBeenCalledWith('1');
  });

  it('triggers parsing when a media project has no frames yet', async () => {
    apiMock.getProjectFrames
      .mockResolvedValueOnce([])
      .mockResolvedValueOnce([
        { id: 11, project_id: 1, frame_index: 0, image_url: '/parsed.jpg', width: 320, height: 240 },
      ]);
    apiMock.parseMedia.mockResolvedValueOnce({ id: 7, status: 'queued', progress: 0 });
    apiMock.getTask.mockResolvedValueOnce({ id: 7, status: 'success', progress: 100, message: '解析完成' });

    render(<VideoWorkspace />);

    await waitFor(() => expect(apiMock.parseMedia).toHaveBeenCalledWith('1'));
    expect(apiMock.getTask).toHaveBeenCalledWith(7);
    await waitFor(() => expect(useStore.getState().frames[0]).toEqual(expect.objectContaining({
      id: '11',
      url: '/parsed.jpg',
    })));
  });

  it('hydrates saved annotations after loading frames', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);
    apiMock.getProjectAnnotations.mockResolvedValueOnce([{ id: 99, frame_id: 10 }]);
    apiMock.annotationToMask.mockReturnValueOnce({
      id: 'annotation-99',
      annotationId: '99',
      frameId: '10',
      saved: true,
      pathData: 'M 0 0 Z',
      label: 'Saved',
      color: '#06b6d4',
    });

    render(<VideoWorkspace />);

    await waitFor(() => expect(useStore.getState().masks).toEqual([
      expect.objectContaining({ id: 'annotation-99', saved: true }),
    ]));
  });

  it('saves pending masks through the archive button', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);
    apiMock.buildAnnotationPayload.mockReturnValueOnce({ project_id: 1, frame_id: 10, mask_data: { polygons: [] } });
    apiMock.saveAnnotation.mockResolvedValueOnce({ id: 5 });

    render(<VideoWorkspace />);
    await waitFor(() => expect(useStore.getState().frames).toHaveLength(1));
    act(() => {
      useStore.setState({
        activeTemplateId: '2',
        masks: [{
          id: 'mask-1',
          frameId: '10',
          pathData: 'M 0 0 Z',
          label: 'AI Mask',
          color: '#06b6d4',
          segmentation: [[0, 0, 10, 0, 10, 10]],
          bbox: [0, 0, 10, 10],
        }],
      });
    });

    fireEvent.click(screen.getByRole('button', { name: '结构化归档保存' }));

    await waitFor(() => expect(apiMock.saveAnnotation).toHaveBeenCalledWith({
      project_id: 1,
      frame_id: 10,
      mask_data: { polygons: [] },
    }));
    expect(apiMock.buildAnnotationPayload).toHaveBeenCalledWith(
      '1',
      expect.objectContaining({ id: 'mask-1' }),
      expect.objectContaining({ id: '10' }),
      '2',
    );
  });

  it('updates dirty saved masks through the archive button', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);
    apiMock.buildAnnotationPayload.mockReturnValueOnce({
      project_id: 1,
      frame_id: 10,
      template_id: 2,
      mask_data: { polygons: [], label: '胆囊' },
    });
    apiMock.updateAnnotation.mockResolvedValueOnce({ id: 99 });

    render(<VideoWorkspace />);
    await waitFor(() => expect(useStore.getState().frames).toHaveLength(1));
    act(() => {
      useStore.setState({
        activeTemplateId: '2',
        masks: [{
          id: 'annotation-99',
          annotationId: '99',
          frameId: '10',
          pathData: 'M 0 0 Z',
          label: '胆囊',
          color: '#ff0000',
          saveStatus: 'dirty',
          segmentation: [[0, 0, 10, 0, 10, 10]],
          bbox: [0, 0, 10, 10],
        }],
      });
    });

    fireEvent.click(screen.getByRole('button', { name: '结构化归档保存' }));

    await waitFor(() => expect(apiMock.updateAnnotation).toHaveBeenCalledWith('99', {
      template_id: 2,
      mask_data: { polygons: [], label: '胆囊' },
      points: undefined,
      bbox: undefined,
    }));
    expect(apiMock.saveAnnotation).not.toHaveBeenCalled();
  });

  it('deletes saved annotations when clearing current-frame masks', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);
    apiMock.deleteAnnotation.mockResolvedValueOnce(undefined);

    render(<VideoWorkspace />);
    await waitFor(() => expect(useStore.getState().frames).toHaveLength(1));
    act(() => {
      useStore.setState({
        masks: [
          {
            id: 'annotation-99',
            annotationId: '99',
            frameId: '10',
            pathData: 'M 0 0 Z',
            label: 'Saved',
            color: '#06b6d4',
            saved: true,
            saveStatus: 'saved',
          },
          {
            id: 'draft-1',
            frameId: '10',
            pathData: 'M 1 1 Z',
            label: 'Draft',
            color: '#ff0000',
          },
        ],
      });
    });

    fireEvent.click(screen.getByRole('button', { name: '清空遮罩' }));

    await waitFor(() => expect(apiMock.deleteAnnotation).toHaveBeenCalledWith('99'));
    expect(useStore.getState().masks).toEqual([]);
  });

  it('auto-saves pending masks before exporting COCO', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);
    apiMock.buildAnnotationPayload.mockReturnValueOnce({ project_id: 1, frame_id: 10, mask_data: { polygons: [] } });
    apiMock.saveAnnotation.mockResolvedValueOnce({ id: 5 });
    apiMock.exportCoco.mockResolvedValueOnce(new Blob(['{}'], { type: 'application/json' }));

    render(<VideoWorkspace />);
    await waitFor(() => expect(useStore.getState().frames).toHaveLength(1));
    act(() => {
      useStore.setState({
        masks: [{
          id: 'mask-1',
          frameId: '10',
          pathData: 'M 0 0 Z',
          label: 'AI Mask',
          color: '#06b6d4',
          segmentation: [[0, 0, 10, 0, 10, 10]],
        }],
      });
    });

    fireEvent.click(screen.getByRole('button', { name: '导出 JSON 标注集' }));

    await waitFor(() => expect(apiMock.saveAnnotation).toHaveBeenCalled());
    expect(apiMock.exportCoco).toHaveBeenCalledWith('1');
  });

  it('auto-saves pending masks before exporting PNG masks', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);
    apiMock.buildAnnotationPayload.mockReturnValueOnce({ project_id: 1, frame_id: 10, mask_data: { polygons: [] } });
    apiMock.saveAnnotation.mockResolvedValueOnce({ id: 5 });
    apiMock.exportMasks.mockResolvedValueOnce(new Blob(['zip'], { type: 'application/zip' }));

    render(<VideoWorkspace />);
    await waitFor(() => expect(useStore.getState().frames).toHaveLength(1));
    act(() => {
      useStore.setState({
        masks: [{
          id: 'mask-1',
          frameId: '10',
          pathData: 'M 0 0 Z',
          label: 'AI Mask',
          color: '#06b6d4',
          segmentation: [[0, 0, 10, 0, 10, 10]],
        }],
      });
    });

    fireEvent.click(screen.getByRole('button', { name: '导出 PNG Mask ZIP' }));

    await waitFor(() => expect(apiMock.saveAnnotation).toHaveBeenCalled());
    expect(apiMock.exportMasks).toHaveBeenCalledWith('1');
  });

  it('imports a GT mask for the current frame and hydrates saved annotations', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
    ]);
    apiMock.importGtMask.mockResolvedValueOnce([{ id: 88, frame_id: 10 }]);
    apiMock.getProjectAnnotations
      .mockResolvedValueOnce([])
      .mockResolvedValueOnce([{ id: 88, frame_id: 10 }]);
    apiMock.annotationToMask.mockReturnValueOnce({
      id: 'annotation-88',
      annotationId: '88',
      frameId: '10',
      saved: true,
      pathData: 'M 0 0 Z',
      label: 'GT Mask',
      color: '#22c55e',
    });

    render(<VideoWorkspace />);
    await waitFor(() => expect(useStore.getState().frames).toHaveLength(1));

    const fileInput = document.querySelector('input[type="file"]') as HTMLInputElement;
    const file = new File(['mask'], 'mask.png', { type: 'image/png' });
    fireEvent.change(fileInput, { target: { files: [file] } });

    await waitFor(() => expect(apiMock.importGtMask).toHaveBeenCalledWith(file, '1', '10'));
    await waitFor(() => expect(useStore.getState().masks).toEqual([
      expect.objectContaining({ id: 'annotation-88', label: 'GT Mask' }),
    ]));
  });

  it('propagates the selected current-frame mask through the backend video tracker', async () => {
    apiMock.getProjectFrames.mockResolvedValueOnce([
      { id: 10, project_id: 1, frame_index: 0, image_url: '/frame.jpg', width: 640, height: 360 },
      { id: 11, project_id: 1, frame_index: 1, image_url: '/frame-1.jpg', width: 640, height: 360 },
    ]);
    apiMock.buildAnnotationPayload.mockReturnValueOnce({
      project_id: 1,
      frame_id: 10,
      template_id: 2,
      mask_data: {
        polygons: [[[0.1, 0.1], [0.3, 0.1], [0.3, 0.3]]],
        label: '胆囊',
        color: '#ff0000',
        class: { id: 'c1', name: '胆囊', color: '#ff0000', zIndex: 20 },
      },
      bbox: [0.1, 0.1, 0.2, 0.2],
    });

    render(<VideoWorkspace />);
    await waitFor(() => expect(useStore.getState().frames).toHaveLength(2));
    act(() => {
      useStore.setState({
        aiModel: 'sam2',
        activeTemplateId: '2',
        selectedMaskIds: ['mask-1'],
        masks: [{
          id: 'mask-1',
          frameId: '10',
          pathData: 'M 0 0 Z',
          label: '胆囊',
          color: '#ff0000',
          segmentation: [[64, 36, 192, 36, 192, 108]],
          bbox: [64, 36, 128, 72],
        }],
      });
    });

    fireEvent.click(screen.getByRole('button', { name: '传播片段' }));

    await waitFor(() => expect(apiMock.propagateMasks).toHaveBeenCalledWith({
      project_id: 1,
      frame_id: 10,
      model: 'sam2',
      direction: 'forward',
      max_frames: 30,
      include_source: false,
      save_annotations: true,
      seed: {
        polygons: [[[0.1, 0.1], [0.3, 0.1], [0.3, 0.3]]],
        bbox: [0.1, 0.1, 0.2, 0.2],
        points: undefined,
        label: '胆囊',
        color: '#ff0000',
        class_metadata: { id: 'c1', name: '胆囊', color: '#ff0000', zIndex: 20 },
        template_id: 2,
      },
    }));
    await waitFor(() => expect(screen.getByText('已传播并保存 2 个区域')).toBeInTheDocument());
  });
});