LocalMind/__tests__/ChatScreen.test.tsx

import { act, fireEvent, render, screen } from '@testing-library/react-native';
import React from 'react';
import ChatScreen from '../app/chat/index';
import * as api from '../services/lmStudioApi';
import { useAppStore } from '../store/useAppStore';

jest.mock('../services/lmStudioApi');
jest.mock('react-native-reanimated', () => {
    const { View } = require('react-native');
    const chain: any = new Proxy(() => chain, { get: () => chain });
    return { __esModule: true, default: { View }, FadeInDown: chain };
});
jest.mock('../components/chat/MessageBubble', () => {
    const { Text } = require('react-native');
    return {
        MessageBubble: ({ message, isLastAssistant }: any) => (
            <Text testID={`bubble-${message.id}`}>{`${message.role}:${message.content}${isLastAssistant ? ':last' : ''}`}</Text>
        ),
    };
});
jest.mock('../components/chat/TypingIndicator', () => {
    const { Text } = require('react-native');
    return { TypingIndicator: () => <Text>typing-indicator</Text> };
});

let inputProps: any;
jest.mock('../components/chat/ChatInput', () => ({
    ChatInput: (props: any) => {
        inputProps = props;
        return null;
    },
}));

const streamMock = api.chatCompletionStream as jest.Mock;
const state = () => useAppStore.getState();

beforeEach(() => {
    jest.useFakeTimers();
    streamMock.mockReset();
    useAppStore.setState(useAppStore.getInitialState(), true);
    useAppStore.setState({ isConnected: true, activeModelId: 'model-a' });
});

afterEach(() => {
    jest.useRealTimers();
});

describe('ChatScreen', () => {
    it('shows the empty state with prompt suggestions when there are no messages', () => {
        render(<ChatScreen />);
        expect(screen.getByText('Start a conversation')).toBeTruthy();
        expect(screen.getByText('Explain recursion in simple terms')).toBeTruthy();
    });

    it('sends a suggestion as a user message and starts streaming', () => {
        render(<ChatScreen />);
        fireEvent.press(screen.getByText('Explain recursion in simple terms'));

        expect(state().messages.map((m) => `${m.role}:${m.content}`)).toEqual([
            'user:Explain recursion in simple terms',
            'assistant:',
        ]);
        expect(streamMock).toHaveBeenCalledTimes(1);
        expect(streamMock.mock.calls[0][1]).toBe('model-a');
    });

    it('hides the empty state and lists the messages once there are some', () => {
        state().addMessage({ role: 'user', content: 'hello' });
        render(<ChatScreen />);
        expect(screen.queryByText('Start a conversation')).toBeNull();
        expect(screen.getByText('user:hello')).toBeTruthy();
    });

    it('shows the offline banner only when the server is not connected', () => {
        const { rerender } = render(<ChatScreen />);
        expect(screen.queryByText(/Server offline/)).toBeNull();

        act(() => {
            state().setIsConnected(false);
        });
        rerender(<ChatScreen />);
        expect(screen.getByText(/Server offline/)).toBeTruthy();
    });

    it('shows the typing indicator only while an empty assistant message is streaming', () => {
        state().addMessage({ role: 'user', content: 'q' });
        const assistantId = state().addMessage({ role: 'assistant', content: '' });
        act(() => {
            state().setIsStreaming(true);
        });
        render(<ChatScreen />);
        expect(screen.getByText('typing-indicator')).toBeTruthy();

        act(() => {
            state().appendToMessage(assistantId, 'partial');
        });
        expect(screen.queryByText('typing-indicator')).toBeNull();
    });

    it('does not show the typing indicator when not streaming', () => {
        state().addMessage({ role: 'assistant', content: '' });
        render(<ChatScreen />);
        expect(screen.queryByText('typing-indicator')).toBeNull();
    });

    it('marks only the last assistant message as regenerable', () => {
        state().addMessage({ role: 'user', content: 'q1' });
        state().addMessage({ role: 'assistant', content: 'a1' });
        state().addMessage({ role: 'user', content: 'q2' });
        state().addMessage({ role: 'assistant', content: 'a2' });
        render(<ChatScreen />);
        expect(screen.getByText('assistant:a1')).toBeTruthy();
        expect(screen.getByText('assistant:a2:last')).toBeTruthy();
    });

    it('sends typed text through the normal path', () => {
        render(<ChatScreen />);
        act(() => {
            inputProps.onSend('typed');
        });
        expect(state().messages[0]).toMatchObject({ role: 'user', content: 'typed' });
        expect(streamMock).toHaveBeenCalledTimes(1);
    });

    describe('editing a message', () => {
        const setup = () => {
            state().addMessage({ role: 'user', content: 'old question' });
            state().addMessage({ role: 'assistant', content: 'old answer' });
            const id = state().messages[0].id;
            render(<ChatScreen />);
            act(() => {
            state().setEditingMessageId(id);
        });
            return id;
        };

        it('passes the edited message text to the input and clears edit mode on cancel', () => {
            setup();
            expect(inputProps.editingText).toBe('old question');
            act(() => {
            inputProps.onCancelEdit();
        });
            expect(state().editingMessageId).toBeNull();
        });

        it('rewrites the message, drops the answers after it and regenerates', () => {
            const id = setup();
            act(() => {
            inputProps.onSend('new question');
        });
            act(() => {
                jest.advanceTimersByTime(100);
            });

            expect(state().editingMessageId).toBeNull();
            expect(state().messages[0]).toMatchObject({ id, content: 'new question', isEdited: true });
            expect(state().messages.map((m) => `${m.role}:${m.content}`)).toEqual([
                'user:new question',
                'assistant:',
            ]);
            expect(streamMock).toHaveBeenCalledTimes(1);
            const [, modelId, apiMessages] = streamMock.mock.calls[0];
            expect(modelId).toBe('model-a');
            expect(apiMessages.slice(1)).toEqual([{ role: 'user', content: 'new question' }]);
            expect(apiMessages[0].role).toBe('system');
        });

        it('edits and explains that a model is needed when none is selected', () => {
            const id = setup();
            act(() => {
                useAppStore.setState({ activeModelId: '' });
            });
            act(() => {
            inputProps.onSend('new question');
        });
            act(() => {
                jest.advanceTimersByTime(100);
            });

            expect(state().messages).toHaveLength(2);
            expect(state().messages[0]).toMatchObject({ id, content: 'new question' });
            expect(state().messages[1].content).toContain('No model selected');
            expect(streamMock).not.toHaveBeenCalled();
        });

        it('lets the user stop an answer generated after an edit', () => {
            const abort = jest.fn();
            streamMock.mockReturnValue({ abort });
            setup();
            act(() => {
            inputProps.onSend('new question');
        });
            act(() => {
                jest.advanceTimersByTime(100);
            });
            act(() => {
            inputProps.onStop();
        });
            expect(abort).toHaveBeenCalledTimes(1);
            expect(state().isStreaming).toBe(false);
        });

        it('saves the edited conversation when the regenerated answer finishes', () => {
            streamMock.mockImplementation((_c, _m, _msgs, onChunk, onDone) => {
                onChunk('fresh answer');
                onDone();
                return { abort: jest.fn() };
            });
            setup();
            act(() => {
            inputProps.onSend('new question');
        });
            act(() => {
                jest.advanceTimersByTime(100);
            });

            expect(state().chatSessions).toHaveLength(1);
            expect(state().chatSessions[0].messages.map((m) => m.content)).toEqual(['new question', 'fresh answer']);
        });
    });
});