openCLI / packages /cli /src /ui /contexts /SessionContext.test.tsx

Upload folder using huggingface_hub

40e575e verified over 1 year ago

7.87 kB

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

import { type MutableRefObject } from 'react';
import { render } from 'ink-testing-library';
import { act } from 'react-dom/test-utils';
import { SessionStatsProvider, useSessionStats } from './SessionContext.js';
import { describe, it, expect, vi } from 'vitest';
import { GenerateContentResponseUsageMetadata } from '@google/genai';

// Mock data that simulates what the Gemini API would return.
const mockMetadata1: GenerateContentResponseUsageMetadata = {
  promptTokenCount: 100,
  candidatesTokenCount: 200,
  totalTokenCount: 300,
  cachedContentTokenCount: 50,
  toolUsePromptTokenCount: 10,
  thoughtsTokenCount: 20,
};

const mockMetadata2: GenerateContentResponseUsageMetadata = {
  promptTokenCount: 10,
  candidatesTokenCount: 20,
  totalTokenCount: 30,
  cachedContentTokenCount: 5,
  toolUsePromptTokenCount: 1,
  thoughtsTokenCount: 2,
};

/**
 * A test harness component that uses the hook and exposes the context value
 * via a mutable ref. This allows us to interact with the context's functions
 * and assert against its state directly in our tests.
 */
const TestHarness = ({
  contextRef,
}: {
  contextRef: MutableRefObject<ReturnType<typeof useSessionStats> | undefined>;
}) => {
  contextRef.current = useSessionStats();
  return null;
};

describe('SessionStatsContext', () => {
  it('should provide the correct initial state', () => {
    const contextRef: MutableRefObject<
      ReturnType<typeof useSessionStats> | undefined
    > = { current: undefined };

    render(
      <SessionStatsProvider>
        <TestHarness contextRef={contextRef} />
      </SessionStatsProvider>,
    );

    const stats = contextRef.current?.stats;

    expect(stats?.sessionStartTime).toBeInstanceOf(Date);
    expect(stats?.currentTurn).toBeDefined();
    expect(stats?.cumulative.turnCount).toBe(0);
    expect(stats?.cumulative.totalTokenCount).toBe(0);
    expect(stats?.cumulative.promptTokenCount).toBe(0);
  });

  it('should increment turnCount when startNewTurn is called', () => {
    const contextRef: MutableRefObject<
      ReturnType<typeof useSessionStats> | undefined
    > = { current: undefined };

    render(
      <SessionStatsProvider>
        <TestHarness contextRef={contextRef} />
      </SessionStatsProvider>,
    );

    act(() => {
      contextRef.current?.startNewTurn();
    });

    const stats = contextRef.current?.stats;
    expect(stats?.currentTurn.totalTokenCount).toBe(0);
    expect(stats?.cumulative.turnCount).toBe(1);
    // Ensure token counts are unaffected
    expect(stats?.cumulative.totalTokenCount).toBe(0);
  });

  it('should aggregate token usage correctly when addUsage is called', () => {
    const contextRef: MutableRefObject<
      ReturnType<typeof useSessionStats> | undefined
    > = { current: undefined };

    render(
      <SessionStatsProvider>
        <TestHarness contextRef={contextRef} />
      </SessionStatsProvider>,
    );

    act(() => {
      contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 123 });
    });

    const stats = contextRef.current?.stats;

    // Check that token counts are updated
    expect(stats?.cumulative.totalTokenCount).toBe(
      mockMetadata1.totalTokenCount ?? 0,
    );
    expect(stats?.cumulative.promptTokenCount).toBe(
      mockMetadata1.promptTokenCount ?? 0,
    );
    expect(stats?.cumulative.apiTimeMs).toBe(123);

    // Check that turn count is NOT incremented
    expect(stats?.cumulative.turnCount).toBe(0);

    // Check that currentTurn is updated
    expect(stats?.currentTurn?.totalTokenCount).toEqual(
      mockMetadata1.totalTokenCount,
    );
    expect(stats?.currentTurn?.apiTimeMs).toBe(123);
  });

  it('should correctly track a full logical turn with multiple API calls', () => {
    const contextRef: MutableRefObject<
      ReturnType<typeof useSessionStats> | undefined
    > = { current: undefined };

    render(
      <SessionStatsProvider>
        <TestHarness contextRef={contextRef} />
      </SessionStatsProvider>,
    );

    // 1. User starts a new turn
    act(() => {
      contextRef.current?.startNewTurn();
    });

    // 2. First API call (e.g., prompt with a tool request)
    act(() => {
      contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
    });

    // 3. Second API call (e.g., sending tool response back)
    act(() => {
      contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
    });

    const stats = contextRef.current?.stats;

    // Turn count should only be 1
    expect(stats?.cumulative.turnCount).toBe(1);

    // --- Check Cumulative Stats ---
    // These fields should be the SUM of both calls
    expect(stats?.cumulative.totalTokenCount).toBe(300 + 30);
    expect(stats?.cumulative.candidatesTokenCount).toBe(200 + 20);
    expect(stats?.cumulative.thoughtsTokenCount).toBe(20 + 2);
    expect(stats?.cumulative.apiTimeMs).toBe(100 + 50);

    // These fields should be the SUM of both calls
    expect(stats?.cumulative.promptTokenCount).toBe(100 + 10);
    expect(stats?.cumulative.cachedContentTokenCount).toBe(50 + 5);
    expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10 + 1);

    // --- Check Current Turn Stats ---
    // All fields should be the SUM of both calls for the turn
    expect(stats?.currentTurn.totalTokenCount).toBe(300 + 30);
    expect(stats?.currentTurn.candidatesTokenCount).toBe(200 + 20);
    expect(stats?.currentTurn.thoughtsTokenCount).toBe(20 + 2);
    expect(stats?.currentTurn.promptTokenCount).toBe(100 + 10);
    expect(stats?.currentTurn.cachedContentTokenCount).toBe(50 + 5);
    expect(stats?.currentTurn.toolUsePromptTokenCount).toBe(10 + 1);
    expect(stats?.currentTurn.apiTimeMs).toBe(100 + 50);
  });

  it('should overwrite currentResponse with each API call', () => {
    const contextRef: MutableRefObject<
      ReturnType<typeof useSessionStats> | undefined
    > = { current: undefined };

    render(
      <SessionStatsProvider>
        <TestHarness contextRef={contextRef} />
      </SessionStatsProvider>,
    );

    // 1. First API call
    act(() => {
      contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
    });

    let stats = contextRef.current?.stats;

    // currentResponse should match the first call
    expect(stats?.currentResponse.totalTokenCount).toBe(300);
    expect(stats?.currentResponse.apiTimeMs).toBe(100);

    // 2. Second API call
    act(() => {
      contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
    });

    stats = contextRef.current?.stats;

    // currentResponse should now match the second call
    expect(stats?.currentResponse.totalTokenCount).toBe(30);
    expect(stats?.currentResponse.apiTimeMs).toBe(50);

    // 3. Start a new turn
    act(() => {
      contextRef.current?.startNewTurn();
    });

    stats = contextRef.current?.stats;

    // currentResponse should be reset
    expect(stats?.currentResponse.totalTokenCount).toBe(0);
    expect(stats?.currentResponse.apiTimeMs).toBe(0);
  });

  it('should throw an error when useSessionStats is used outside of a provider', () => {
    // Suppress the expected console error during this test.
    const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});

    const contextRef = { current: undefined };

    // We expect rendering to fail, which React will catch and log as an error.
    render(<TestHarness contextRef={contextRef} />);

    // Assert that the first argument of the first call to console.error
    // contains the expected message. This is more robust than checking
    // the exact arguments, which can be affected by React/JSDOM internals.
    expect(errorSpy.mock.calls[0][0]).toContain(
      'useSessionStats must be used within a SessionStatsProvider',
    );

    errorSpy.mockRestore();
  });
});

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

	import { type MutableRefObject } from 'react';
	import { render } from 'ink-testing-library';
	import { act } from 'react-dom/test-utils';
	import { SessionStatsProvider, useSessionStats } from './SessionContext.js';
	import { describe, it, expect, vi } from 'vitest';
	import { GenerateContentResponseUsageMetadata } from '@google/genai';

	// Mock data that simulates what the Gemini API would return.
	const mockMetadata1: GenerateContentResponseUsageMetadata = {
	promptTokenCount: 100,
	candidatesTokenCount: 200,
	totalTokenCount: 300,
	cachedContentTokenCount: 50,
	toolUsePromptTokenCount: 10,
	thoughtsTokenCount: 20,
	};

	const mockMetadata2: GenerateContentResponseUsageMetadata = {
	promptTokenCount: 10,
	candidatesTokenCount: 20,
	totalTokenCount: 30,
	cachedContentTokenCount: 5,
	toolUsePromptTokenCount: 1,
	thoughtsTokenCount: 2,
	};

	/**
	* A test harness component that uses the hook and exposes the context value
	* via a mutable ref. This allows us to interact with the context's functions
	* and assert against its state directly in our tests.
	*/
	const TestHarness = ({
	contextRef,
	}: {
	contextRef: MutableRefObject<ReturnType<typeof useSessionStats> \| undefined>;
	}) => {
	contextRef.current = useSessionStats();
	return null;
	};

	describe('SessionStatsContext', () => {
	it('should provide the correct initial state', () => {
	const contextRef: MutableRefObject<
	ReturnType<typeof useSessionStats> \| undefined
	> = { current: undefined };

	render(
	<SessionStatsProvider>
	<TestHarness contextRef={contextRef} />
	</SessionStatsProvider>,
	);

	const stats = contextRef.current?.stats;

	expect(stats?.sessionStartTime).toBeInstanceOf(Date);
	expect(stats?.currentTurn).toBeDefined();
	expect(stats?.cumulative.turnCount).toBe(0);
	expect(stats?.cumulative.totalTokenCount).toBe(0);
	expect(stats?.cumulative.promptTokenCount).toBe(0);
	});

	it('should increment turnCount when startNewTurn is called', () => {
	const contextRef: MutableRefObject<
	ReturnType<typeof useSessionStats> \| undefined
	> = { current: undefined };

	render(
	<SessionStatsProvider>
	<TestHarness contextRef={contextRef} />
	</SessionStatsProvider>,
	);

	act(() => {
	contextRef.current?.startNewTurn();
	});

	const stats = contextRef.current?.stats;
	expect(stats?.currentTurn.totalTokenCount).toBe(0);
	expect(stats?.cumulative.turnCount).toBe(1);
	// Ensure token counts are unaffected
	expect(stats?.cumulative.totalTokenCount).toBe(0);
	});

	it('should aggregate token usage correctly when addUsage is called', () => {
	const contextRef: MutableRefObject<
	ReturnType<typeof useSessionStats> \| undefined
	> = { current: undefined };

	render(
	<SessionStatsProvider>
	<TestHarness contextRef={contextRef} />
	</SessionStatsProvider>,
	);

	act(() => {
	contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 123 });
	});

	const stats = contextRef.current?.stats;

	// Check that token counts are updated
	expect(stats?.cumulative.totalTokenCount).toBe(
	mockMetadata1.totalTokenCount ?? 0,
	);
	expect(stats?.cumulative.promptTokenCount).toBe(
	mockMetadata1.promptTokenCount ?? 0,
	);
	expect(stats?.cumulative.apiTimeMs).toBe(123);

	// Check that turn count is NOT incremented
	expect(stats?.cumulative.turnCount).toBe(0);

	// Check that currentTurn is updated
	expect(stats?.currentTurn?.totalTokenCount).toEqual(
	mockMetadata1.totalTokenCount,
	);
	expect(stats?.currentTurn?.apiTimeMs).toBe(123);
	});

	it('should correctly track a full logical turn with multiple API calls', () => {
	const contextRef: MutableRefObject<
	ReturnType<typeof useSessionStats> \| undefined
	> = { current: undefined };

	render(
	<SessionStatsProvider>
	<TestHarness contextRef={contextRef} />
	</SessionStatsProvider>,
	);

	// 1. User starts a new turn
	act(() => {
	contextRef.current?.startNewTurn();
	});

	// 2. First API call (e.g., prompt with a tool request)
	act(() => {
	contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
	});

	// 3. Second API call (e.g., sending tool response back)
	act(() => {
	contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
	});

	const stats = contextRef.current?.stats;

	// Turn count should only be 1
	expect(stats?.cumulative.turnCount).toBe(1);

	// --- Check Cumulative Stats ---
	// These fields should be the SUM of both calls
	expect(stats?.cumulative.totalTokenCount).toBe(300 + 30);
	expect(stats?.cumulative.candidatesTokenCount).toBe(200 + 20);
	expect(stats?.cumulative.thoughtsTokenCount).toBe(20 + 2);
	expect(stats?.cumulative.apiTimeMs).toBe(100 + 50);

	// These fields should be the SUM of both calls
	expect(stats?.cumulative.promptTokenCount).toBe(100 + 10);
	expect(stats?.cumulative.cachedContentTokenCount).toBe(50 + 5);
	expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10 + 1);

	// --- Check Current Turn Stats ---
	// All fields should be the SUM of both calls for the turn
	expect(stats?.currentTurn.totalTokenCount).toBe(300 + 30);
	expect(stats?.currentTurn.candidatesTokenCount).toBe(200 + 20);
	expect(stats?.currentTurn.thoughtsTokenCount).toBe(20 + 2);
	expect(stats?.currentTurn.promptTokenCount).toBe(100 + 10);
	expect(stats?.currentTurn.cachedContentTokenCount).toBe(50 + 5);
	expect(stats?.currentTurn.toolUsePromptTokenCount).toBe(10 + 1);
	expect(stats?.currentTurn.apiTimeMs).toBe(100 + 50);
	});

	it('should overwrite currentResponse with each API call', () => {
	const contextRef: MutableRefObject<
	ReturnType<typeof useSessionStats> \| undefined
	> = { current: undefined };

	render(
	<SessionStatsProvider>
	<TestHarness contextRef={contextRef} />
	</SessionStatsProvider>,
	);

	// 1. First API call
	act(() => {
	contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 });
	});

	let stats = contextRef.current?.stats;

	// currentResponse should match the first call
	expect(stats?.currentResponse.totalTokenCount).toBe(300);
	expect(stats?.currentResponse.apiTimeMs).toBe(100);

	// 2. Second API call
	act(() => {
	contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 });
	});

	stats = contextRef.current?.stats;

	// currentResponse should now match the second call
	expect(stats?.currentResponse.totalTokenCount).toBe(30);
	expect(stats?.currentResponse.apiTimeMs).toBe(50);

	// 3. Start a new turn
	act(() => {
	contextRef.current?.startNewTurn();
	});

	stats = contextRef.current?.stats;

	// currentResponse should be reset
	expect(stats?.currentResponse.totalTokenCount).toBe(0);
	expect(stats?.currentResponse.apiTimeMs).toBe(0);
	});

	it('should throw an error when useSessionStats is used outside of a provider', () => {
	// Suppress the expected console error during this test.
	const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});

	const contextRef = { current: undefined };

	// We expect rendering to fail, which React will catch and log as an error.
	render(<TestHarness contextRef={contextRef} />);

	// Assert that the first argument of the first call to console.error
	// contains the expected message. This is more robust than checking
	// the exact arguments, which can be affected by React/JSDOM internals.
	expect(errorSpy.mock.calls[0][0]).toContain(
	'useSessionStats must be used within a SessionStatsProvider',
	);

	errorSpy.mockRestore();
	});
	});