Genarrative/src/components/image-editor/EditorAgentConversation/ToolCallView.test.tsx

/* @vitest-environment jsdom */

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

import type { EditorAgentToolCall } from '@/packages/shared/src/contracts';

import ToolCallView from './ToolCallView.tsx';

const getExternalGenerationJobStatusMock = vi.hoisted(() => vi.fn());

vi.mock('@/src/services/external-generation', () => ({
  getExternalGenerationJobStatus: getExternalGenerationJobStatusMock,
}));

function createDeferred<T>() {
  let resolve!: (value: T) => void;
  const promise = new Promise<T>((nextResolve) => {
    resolve = nextResolve;
  });
  return { promise, resolve };
}

function createToolCall(
  overrides: Partial<EditorAgentToolCall> = {},
): EditorAgentToolCall {
  return {
    toolName: 'generate-image',
    status: 'not_completed',
    args: {},
    displayArgs: {
      stringArgs: [],
      imageArgs: [],
      extras: { priceMudPoints: 1 },
    },
    externalJobId: 'job-a',
    images: [],
    error: null,
    ...overrides,
  };
}

function createCompletedJobResponse(jobId: string) {
  return {
    job: {
      operationId: jobId,
      status: 'completed' as const,
      phaseLabel: '已完成',
      phaseDetail: '生成完成',
      progress: 100,
      error: null,
      updatedAtMicros: 1,
    },
  };
}

describe('ToolCallView', () => {
  beforeEach(() => {
    getExternalGenerationJobStatusMock.mockReset();
  });

  it.each([
    ['failed', '失败', '生成失败'],
    ['cancelled', '已取消', null],
  ] as const)(
    'keeps the server %s state when the pending response arrives late',
    async (status, statusLabel, error) => {
      const pendingResponse = createDeferred<
        ReturnType<typeof createCompletedJobResponse>
      >();
      getExternalGenerationJobStatusMock.mockReturnValueOnce(
        pendingResponse.promise,
      );
      const onJobCompleted = vi.fn();
      const { rerender } = render(
        <ToolCallView
          toolCall={createToolCall()}
          onJobCompleted={onJobCompleted}
        />,
      );

      await waitFor(() => {
        expect(getExternalGenerationJobStatusMock).toHaveBeenCalledWith(
          'job-a',
        );
      });

      rerender(
        <ToolCallView
          toolCall={createToolCall({ status, error })}
          onJobCompleted={onJobCompleted}
        />,
      );
      expect(await screen.findByText(statusLabel)).toBeTruthy();

      await act(async () => {
        pendingResponse.resolve(createCompletedJobResponse('job-a'));
        await pendingResponse.promise;
      });

      expect(screen.getByText(statusLabel)).toBeTruthy();
      expect(screen.queryByText('已完成')).toBeNull();
      expect(onJobCompleted).not.toHaveBeenCalled();
    },
  );

  it('ignores job A after switching to job B and completes job B once', async () => {
    const jobAResponse = createDeferred<
      ReturnType<typeof createCompletedJobResponse>
    >();
    const jobBResponse = createDeferred<
      ReturnType<typeof createCompletedJobResponse>
    >();
    getExternalGenerationJobStatusMock.mockImplementation((jobId: string) =>
      jobId === 'job-a' ? jobAResponse.promise : jobBResponse.promise,
    );
    const onJobCompleted = vi.fn();
    const { rerender } = render(
      <ToolCallView
        toolCall={createToolCall()}
        onJobCompleted={onJobCompleted}
      />,
    );

    await waitFor(() => {
      expect(getExternalGenerationJobStatusMock).toHaveBeenCalledWith('job-a');
    });

    rerender(
      <ToolCallView
        toolCall={createToolCall({ externalJobId: 'job-b' })}
        onJobCompleted={onJobCompleted}
      />,
    );
    await waitFor(() => {
      expect(getExternalGenerationJobStatusMock).toHaveBeenCalledWith('job-b');
    });

    await act(async () => {
      jobAResponse.resolve(createCompletedJobResponse('job-a'));
      await jobAResponse.promise;
    });
    expect(screen.getByText('执行中')).toBeTruthy();
    expect(onJobCompleted).not.toHaveBeenCalled();

    await act(async () => {
      jobBResponse.resolve(createCompletedJobResponse('job-b'));
      await jobBResponse.promise;
    });

    expect(await screen.findByText('已完成')).toBeTruthy();
    expect(onJobCompleted).toHaveBeenCalledTimes(1);
  });
});