qwen-code/packages/cli/src/ui/hooks/useLoadingIndicator.test.ts

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

import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { renderHook, act } from '@testing-library/react';
import { useLoadingIndicator } from './useLoadingIndicator.js';
import { StreamingState } from '../types.js';
import { PHRASE_CHANGE_INTERVAL_MS } from './usePhraseCycler.js';
import * as i18n from '../../i18n/index.js';

const MOCK_WITTY_PHRASES = ['Phrase 1', 'Phrase 2', 'Phrase 3'];

describe('useLoadingIndicator', () => {
  beforeEach(() => {
    vi.useFakeTimers();
    vi.spyOn(i18n, 'ta').mockReturnValue(MOCK_WITTY_PHRASES);
    vi.spyOn(i18n, 't').mockImplementation((key) => key);
  });

  afterEach(() => {
    vi.useRealTimers(); // Restore real timers after each test
    act(() => vi.runOnlyPendingTimers);
    vi.restoreAllMocks();
  });

  it('should initialize with default values when Idle', () => {
    const { result } = renderHook(() =>
      useLoadingIndicator(StreamingState.Idle),
    );
    expect(result.current.elapsedTime).toBe(0);
    expect(MOCK_WITTY_PHRASES).toContain(result.current.currentLoadingPhrase);
  });

  it('should reflect values when Responding', async () => {
    const { result } = renderHook(() =>
      useLoadingIndicator(StreamingState.Responding),
    );

    // Initial state before timers advance
    expect(result.current.elapsedTime).toBe(0);
    expect(MOCK_WITTY_PHRASES).toContain(result.current.currentLoadingPhrase);

    await act(async () => {
      await vi.advanceTimersByTimeAsync(PHRASE_CHANGE_INTERVAL_MS + 1);
    });

    // Phrase should cycle if PHRASE_CHANGE_INTERVAL_MS has passed
    expect(MOCK_WITTY_PHRASES).toContain(result.current.currentLoadingPhrase);
  });

  it('should show waiting phrase and retain elapsedTime when WaitingForConfirmation', async () => {
    const { result, rerender } = renderHook(
      ({ streamingState }) => useLoadingIndicator(streamingState),
      { initialProps: { streamingState: StreamingState.Responding } },
    );

    await act(async () => {
      await vi.advanceTimersByTimeAsync(60000);
    });
    expect(result.current.elapsedTime).toBe(60);

    act(() => {
      rerender({ streamingState: StreamingState.WaitingForConfirmation });
    });

    expect(result.current.currentLoadingPhrase).toBe(
      'Waiting for user confirmation...',
    );
    expect(result.current.elapsedTime).toBe(60); // Elapsed time should be retained

    // Timer should not advance further
    await act(async () => {
      await vi.advanceTimersByTimeAsync(2000);
    });
    expect(result.current.elapsedTime).toBe(60);
  });

  it('should reset elapsedTime and use a witty phrase when transitioning from WaitingForConfirmation to Responding', async () => {
    const { result, rerender } = renderHook(
      ({ streamingState }) => useLoadingIndicator(streamingState),
      { initialProps: { streamingState: StreamingState.Responding } },
    );

    await act(async () => {
      await vi.advanceTimersByTimeAsync(5000); // 5s
    });
    expect(result.current.elapsedTime).toBe(5);

    act(() => {
      rerender({ streamingState: StreamingState.WaitingForConfirmation });
    });
    expect(result.current.elapsedTime).toBe(5);
    expect(result.current.currentLoadingPhrase).toBe(
      'Waiting for user confirmation...',
    );

    act(() => {
      rerender({ streamingState: StreamingState.Responding });
    });
    expect(result.current.elapsedTime).toBe(0); // Should reset
    expect(MOCK_WITTY_PHRASES).toContain(result.current.currentLoadingPhrase);

    await act(async () => {
      await vi.advanceTimersByTimeAsync(1000);
    });
    expect(result.current.elapsedTime).toBe(1);
  });

  it('should reset timer and phrase when streamingState changes from Responding to Idle', async () => {
    const { result, rerender } = renderHook(
      ({ streamingState }) => useLoadingIndicator(streamingState),
      { initialProps: { streamingState: StreamingState.Responding } },
    );

    await act(async () => {
      await vi.advanceTimersByTimeAsync(10000); // 10s
    });
    expect(result.current.elapsedTime).toBe(10);

    act(() => {
      rerender({ streamingState: StreamingState.Idle });
    });

    expect(result.current.elapsedTime).toBe(0);
    expect(MOCK_WITTY_PHRASES).toContain(result.current.currentLoadingPhrase);

    // Timer should not advance
    await act(async () => {
      await vi.advanceTimersByTimeAsync(2000);
    });
    expect(result.current.elapsedTime).toBe(0);
  });

  describe('token tracking', () => {
    it('should capture token snapshot when task starts', () => {
      const { result, rerender } = renderHook(
        ({ streamingState, currentCandidatesTokens }) =>
          useLoadingIndicator(
            streamingState,
            undefined,
            currentCandidatesTokens,
          ),
        {
          initialProps: {
            streamingState: StreamingState.Idle,
            currentCandidatesTokens: 100,
          },
        },
      );

      expect(result.current.taskStartTokens).toBe(0);

      act(() => {
        rerender({
          streamingState: StreamingState.Responding,
          currentCandidatesTokens: 100,
        });
      });

      expect(result.current.taskStartTokens).toBe(100);
    });

    it('should reset token snapshot when transitioning from Responding to Idle', async () => {
      const { result, rerender } = renderHook(
        ({ streamingState, currentCandidatesTokens }) =>
          useLoadingIndicator(
            streamingState,
            undefined,
            currentCandidatesTokens,
          ),
        {
          initialProps: {
            streamingState: StreamingState.Idle,
            currentCandidatesTokens: 0,
          },
        },
      );

      act(() => {
        rerender({
          streamingState: StreamingState.Responding,
          currentCandidatesTokens: 0,
        });
      });
      expect(result.current.taskStartTokens).toBe(0);

      await act(async () => {
        await vi.advanceTimersByTimeAsync(1000);
        rerender({
          streamingState: StreamingState.Responding,
          currentCandidatesTokens: 500,
        });
      });

      act(() => {
        rerender({
          streamingState: StreamingState.Idle,
          currentCandidatesTokens: 500,
        });
      });

      expect(result.current.taskStartTokens).toBe(0);
    });

    it('should reset token snapshot when transitioning from WaitingForConfirmation to Responding', async () => {
      const { result, rerender } = renderHook(
        ({ streamingState, currentCandidatesTokens }) =>
          useLoadingIndicator(
            streamingState,
            undefined,
            currentCandidatesTokens,
          ),
        {
          initialProps: {
            streamingState: StreamingState.Responding,
            currentCandidatesTokens: 100,
          },
        },
      );

      expect(result.current.taskStartTokens).toBe(100);

      await act(async () => {
        await vi.advanceTimersByTimeAsync(5000);
        rerender({
          streamingState: StreamingState.Responding,
          currentCandidatesTokens: 500,
        });
      });

      act(() => {
        rerender({
          streamingState: StreamingState.WaitingForConfirmation,
          currentCandidatesTokens: 500,
        });
      });

      act(() => {
        rerender({
          streamingState: StreamingState.Responding,
          currentCandidatesTokens: 500,
        });
      });

      expect(result.current.taskStartTokens).toBe(500);
    });
  });
});