packages/cli/src/ui/components/BackgroundTaskDisplay.test.tsx

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { render } from '../../test-utils/render.js';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { BackgroundTaskDisplay } from './BackgroundTaskDisplay.js';
import { type BackgroundTask } from '../hooks/useExecutionLifecycle.js';
import { ShellExecutionService } from '@google/gemini-cli-core';
import { act } from 'react';
import { type Key, type KeypressHandler } from '../contexts/KeypressContext.js';
import { ScrollProvider } from '../contexts/ScrollProvider.js';
import { Box } from 'ink';

// Mock dependencies
const mockDismissBackgroundTask = vi.fn();
const mockSetActiveBackgroundTaskPid = vi.fn();
const mockSetIsBackgroundTaskListOpen = vi.fn();

vi.mock('../contexts/UIActionsContext.js', () => ({
  useUIActions: () => ({
    dismissBackgroundTask: mockDismissBackgroundTask,
    setActiveBackgroundTaskPid: mockSetActiveBackgroundTaskPid,
    setIsBackgroundTaskListOpen: mockSetIsBackgroundTaskListOpen,
  }),
}));

vi.mock('@google/gemini-cli-core', async (importOriginal) => {
  const actual =
    await importOriginal<typeof import('@google/gemini-cli-core')>();
  return {
    ...actual,
    ShellExecutionService: {
      resizePty: vi.fn(),
      subscribe: vi.fn(() => vi.fn()),
      getLogFilePath: vi.fn(
        (pid) => `~/.gemini/tmp/background-processes/background-${pid}.log`,
      ),
      getLogDir: vi.fn(() => '~/.gemini/tmp/background-processes'),
    },
  };
});

// Mock AnsiOutputText since it's a complex component
vi.mock('./AnsiOutput.js', () => ({
  AnsiOutputText: ({ data }: { data: string | unknown }) => {
    if (typeof data === 'string') return <>{data}</>;
    // Simple serialization for object data
    return <>{JSON.stringify(data)}</>;
  },
}));

// Mock useKeypress
let keypressHandlers: Array<{ handler: KeypressHandler; isActive: boolean }> =
  [];
vi.mock('../hooks/useKeypress.js', () => ({
  useKeypress: vi.fn((handler, { isActive }) => {
    keypressHandlers.push({ handler, isActive });
  }),
}));

const simulateKey = (key: Partial<Key>) => {
  const fullKey: Key = createMockKey(key);
  keypressHandlers.forEach(({ handler, isActive }) => {
    if (isActive) {
      handler(fullKey);
    }
  });
};

vi.mock('../contexts/MouseContext.js', () => ({
  useMouseContext: vi.fn(() => ({
    subscribe: vi.fn(),
    unsubscribe: vi.fn(),
  })),
  useMouse: vi.fn(),
}));

// Mock ScrollableList
vi.mock('./shared/ScrollableList.js', () => ({
  SCROLL_TO_ITEM_END: 999999,
  ScrollableList: vi.fn(
    ({
      data,
      renderItem,
    }: {
      data: BackgroundTask[];
      renderItem: (props: {
        item: BackgroundTask;
        index: number;
      }) => React.ReactNode;
    }) => (
      <Box flexDirection="column">
        {data.map((item: BackgroundTask, index: number) => (
          <Box key={index}>{renderItem({ item, index })}</Box>
        ))}
      </Box>
    ),
  ),
}));

afterEach(() => {
  vi.restoreAllMocks();
});

const createMockKey = (overrides: Partial<Key>): Key => ({
  name: '',
  ctrl: false,
  alt: false,
  cmd: false,
  shift: false,
  insertable: false,
  sequence: '',
  ...overrides,
});

describe('<BackgroundTaskDisplay />', () => {
  const mockShells = new Map<number, BackgroundTask>();
  const shell1: BackgroundTask = {
    pid: 1001,
    command: 'npm start',
    output: 'Starting server...',
    isBinary: false,
    binaryBytesReceived: 0,
    status: 'running',
  };
  const shell2: BackgroundTask = {
    pid: 1002,
    command: 'tail -f log.txt',
    output: 'Log entry 1',
    isBinary: false,
    binaryBytesReceived: 0,
    status: 'running',
  };

  beforeEach(() => {
    vi.clearAllMocks();
    mockShells.clear();
    mockShells.set(shell1.pid, shell1);
    mockShells.set(shell2.pid, shell2);
    keypressHandlers = [];
  });

  it('renders the output of the active shell', async () => {
    const width = 80;
    const { lastFrame, unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={false}
          isListOpenProp={false}
        />
      </ScrollProvider>,
      width,
    );

    expect(lastFrame()).toMatchSnapshot();
    unmount();
  });

  it('renders tabs for multiple shells', async () => {
    const width = 100;
    const { lastFrame, unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={false}
          isListOpenProp={false}
        />
      </ScrollProvider>,
      width,
    );

    expect(lastFrame()).toMatchSnapshot();
    unmount();
  });

  it('highlights the focused state', async () => {
    const width = 80;
    const { lastFrame, unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={true}
          isListOpenProp={false}
        />
      </ScrollProvider>,
      width,
    );

    expect(lastFrame()).toMatchSnapshot();
    unmount();
  });

  it('resizes the PTY on mount and when dimensions change', async () => {
    const width = 80;
    const { rerender, unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={false}
          isListOpenProp={false}
        />
      </ScrollProvider>,
      width,
    );

    expect(ShellExecutionService.resizePty).toHaveBeenCalledWith(
      shell1.pid,
      76,
      20,
    );

    rerender(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={100}
          height={30}
          isFocused={false}
          isListOpenProp={false}
        />
      </ScrollProvider>,
    );

    expect(ShellExecutionService.resizePty).toHaveBeenCalledWith(
      shell1.pid,
      96,
      26,
    );
    unmount();
  });

  it('renders the process list when isListOpenProp is true', async () => {
    const width = 80;
    const { lastFrame, unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={true}
          isListOpenProp={true}
        />
      </ScrollProvider>,
      width,
    );

    expect(lastFrame()).toMatchSnapshot();
    unmount();
  });

  it('selects the current process and closes the list when Ctrl+L is pressed in list view', async () => {
    const width = 80;
    const { unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={true}
          isListOpenProp={true}
        />
      </ScrollProvider>,
      width,
    );

    // Simulate down arrow to select the second process (handled by RadioButtonSelect)
    await act(async () => {
      simulateKey({ name: 'down' });
    });

    // Simulate Ctrl+L (handled by BackgroundTaskDisplay)
    await act(async () => {
      simulateKey({ name: 'l', ctrl: true });
    });

    expect(mockSetActiveBackgroundTaskPid).toHaveBeenCalledWith(shell2.pid);
    expect(mockSetIsBackgroundTaskListOpen).toHaveBeenCalledWith(false);
    unmount();
  });

  it('kills the highlighted process when Ctrl+K is pressed in list view', async () => {
    const width = 80;
    const { unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={true}
          isListOpenProp={true}
        />
      </ScrollProvider>,
      width,
    );

    // Initial state: shell1 (active) is highlighted

    // Move to shell2
    await act(async () => {
      simulateKey({ name: 'down' });
    });

    // Press Ctrl+K
    await act(async () => {
      simulateKey({ name: 'k', ctrl: true });
    });

    expect(mockDismissBackgroundTask).toHaveBeenCalledWith(shell2.pid);
    unmount();
  });

  it('kills the active process when Ctrl+K is pressed in output view', async () => {
    const width = 80;
    const { unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell1.pid}
          width={width}
          height={24}
          isFocused={true}
          isListOpenProp={false}
        />
      </ScrollProvider>,
      width,
    );

    await act(async () => {
      simulateKey({ name: 'k', ctrl: true });
    });

    expect(mockDismissBackgroundTask).toHaveBeenCalledWith(shell1.pid);
    unmount();
  });

  it('scrolls to active shell when list opens', async () => {
    // shell2 is active
    const width = 80;
    const { lastFrame, unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={shell2.pid}
          width={width}
          height={24}
          isFocused={true}
          isListOpenProp={true}
        />
      </ScrollProvider>,
      width,
    );

    expect(lastFrame()).toMatchSnapshot();
    unmount();
  });

  it('keeps exit code status color even when selected', async () => {
    const exitedShell: BackgroundTask = {
      pid: 1003,
      command: 'exit 0',
      output: '',
      isBinary: false,
      binaryBytesReceived: 0,
      status: 'exited',
      exitCode: 0,
    };
    mockShells.set(exitedShell.pid, exitedShell);

    const width = 80;
    const { lastFrame, unmount } = await render(
      <ScrollProvider>
        <BackgroundTaskDisplay
          shells={mockShells}
          activePid={exitedShell.pid}
          width={width}
          height={24}
          isFocused={true}
          isListOpenProp={true}
        />
      </ScrollProvider>,
      width,
    );

    expect(lastFrame()).toMatchSnapshot();
    unmount();
  });
});