life-echo/app-expo/tests/features/conversation/entry-warmup.test.ts

import { QueryClient } from '@tanstack/react-query';

import {
  prefetchConversationMessages,
  warmupConversationOpening,
} from '@/features/conversation/entry-warmup';
import { conversationKeys } from '@/features/conversation/query-keys';
import type { MessageItem } from '@/features/conversation/types';

const mockLoadMessages = jest.fn();
const mockRegisterPreparedRealtimeSession = jest.fn();
let mockConnectImpl:
  | ((options: {
      conversationId: string;
      queryClient: QueryClient;
    }) => Promise<void> | void)
  | null = null;
const mockSessions: Array<{
  attachUiCallbacks: jest.Mock;
  connect: jest.Mock;
  dispose: jest.Mock;
}> = [];

jest.mock('@/features/conversation/conversation-messages-repository', () => ({
  conversationMessagesRepository: {
    loadMessages: (conversationId: string) => mockLoadMessages(conversationId),
  },
}));

jest.mock('@/features/conversation/prepared-session-registry', () => ({
  registerPreparedRealtimeSession: (conversationId: string, session: unknown) =>
    mockRegisterPreparedRealtimeSession(conversationId, session),
}));

jest.mock('@/features/conversation/realtime-session', () => ({
  RealtimeSession: jest.fn().mockImplementation((options) => {
    const session = {
      attachUiCallbacks: jest.fn(),
      connect: jest.fn(async () => {
        await mockConnectImpl?.(options);
      }),
      dispose: jest.fn(),
    };
    mockSessions.push(session);
    return session;
  }),
}));

function createQueryClient(): QueryClient {
  return new QueryClient({
    defaultOptions: {
      queries: { retry: false, gcTime: Infinity },
      mutations: { retry: false },
    },
  });
}

function assistantMessage(id = 'assistant-1'): MessageItem {
  return {
    id,
    conversationId: 'conv-1',
    content: '你好，今天想聊哪段回忆？',
    senderType: 'assistant',
    timestamp: 1,
    messageType: 'text',
  };
}

describe('conversation entry warmup', () => {
  let queryClient: QueryClient;

  beforeEach(() => {
    queryClient = createQueryClient();
    mockLoadMessages.mockReset();
    mockRegisterPreparedRealtimeSession.mockReset();
    mockConnectImpl = null;
    mockSessions.length = 0;
  });

  afterEach(async () => {
    await queryClient.cancelQueries();
    queryClient.clear();
  });

  test('prefetches messages without throwing on load failure', async () => {
    mockLoadMessages.mockRejectedValueOnce(new Error('network down'));

    await expect(
      prefetchConversationMessages(queryClient, 'conv-1'),
    ).resolves.toBeUndefined();
  });

  test('uses refreshed history and skips websocket when opening is already cached', async () => {
    const existing = assistantMessage();
    mockLoadMessages.mockResolvedValueOnce([existing]);

    await warmupConversationOpening(queryClient, 'conv-1');

    expect(mockLoadMessages).toHaveBeenCalledWith('conv-1');
    expect(mockSessions).toHaveLength(0);
    expect(
      queryClient.getQueryData(conversationKeys.messages('conv-1')),
    ).toEqual([existing]);
  });

  test('connects websocket and registers prepared session after opening arrives', async () => {
    const opened = assistantMessage();
    mockLoadMessages.mockResolvedValueOnce([]).mockResolvedValueOnce([opened]);
    mockConnectImpl = ({ conversationId, queryClient }) => {
      queryClient.setQueryData(conversationKeys.messages(conversationId), [
        opened,
      ]);
    };

    await warmupConversationOpening(queryClient, 'conv-1');

    expect(mockSessions).toHaveLength(1);
    expect(mockSessions[0]?.attachUiCallbacks).toHaveBeenCalled();
    expect(mockSessions[0]?.connect).toHaveBeenCalled();
    expect(mockSessions[0]?.dispose).not.toHaveBeenCalled();
    expect(mockRegisterPreparedRealtimeSession).toHaveBeenCalledWith(
      'conv-1',
      mockSessions[0],
    );
  });
});