refactor(chat): split ChatScreen and reuse regenerateResponse after editing a message
The edit path duplicated the streaming flow with inline require() calls; it now reuses regenerateResponse, which also fixes: - edited conversations were never saved to History - the answer generated after an edit could not be stopped (no abort handle kept) Stopping a stream now saves the partial conversation, like finishing or failing does. Adds characterization tests for ChatScreen and useLMStudio.
This commit is contained in:
evandro committed
1 parent
808cb1610b
commit
cb6eda0562
8 files changed
+678
-176
No files matched your search
@@ -0,0 +1,220 @@
|
||||
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 without regenerating when no model is selected', () => {
|
||||
const id = setup();
|
||||
act(() => {
|
||||
useAppStore.setState({ activeModelId: '' });
|
||||
});
|
||||
act(() => {
|
||||
inputProps.onSend('new question');
|
||||
});
|
||||
act(() => {
|
||||
jest.advanceTimersByTime(100);
|
||||
});
|
||||
|
||||
expect(state().messages).toHaveLength(1);
|
||||
expect(state().messages[0]).toMatchObject({ id, content: 'new question' });
|
||||
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']);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,234 @@
|
||||
import { act, renderHook } from '@testing-library/react-native';
|
||||
import { useLMStudio } from '../hooks/useLMStudio';
|
||||
import * as api from '../services/lmStudioApi';
|
||||
import { useAppStore } from '../store/useAppStore';
|
||||
|
||||
jest.mock('../services/lmStudioApi');
|
||||
|
||||
const mockApi = api as jest.Mocked<typeof api>;
|
||||
const state = () => useAppStore.getState();
|
||||
|
||||
type StreamArgs = Parameters<typeof api.chatCompletionStream>;
|
||||
const lastStreamCall = () => mockApi.chatCompletionStream.mock.calls.at(-1) as StreamArgs;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.resetAllMocks();
|
||||
useAppStore.setState(useAppStore.getInitialState(), true);
|
||||
useAppStore.setState({ activeModelId: 'model-a' });
|
||||
mockApi.chatCompletionStream.mockReturnValue({ abort: jest.fn() });
|
||||
});
|
||||
|
||||
describe('checkConnection', () => {
|
||||
it('stores the result of the connection test', async () => {
|
||||
mockApi.testConnection.mockResolvedValue(true);
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
|
||||
await act(async () => {
|
||||
await expect(result.current.checkConnection()).resolves.toBe(true);
|
||||
});
|
||||
expect(state().isConnected).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('fetchModels', () => {
|
||||
const models = [
|
||||
{ id: 'first', object: 'model', owned_by: 'x' },
|
||||
{ id: 'second', object: 'model', owned_by: 'x' },
|
||||
];
|
||||
|
||||
it('stores the models, marks the server connected and selects the first model if none is active', async () => {
|
||||
useAppStore.setState({ activeModelId: '' });
|
||||
mockApi.getModels.mockResolvedValue(models);
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
|
||||
await act(async () => {
|
||||
await result.current.fetchModels();
|
||||
});
|
||||
expect(state().availableModels).toEqual(models);
|
||||
expect(state().isConnected).toBe(true);
|
||||
expect(state().activeModelId).toBe('first');
|
||||
});
|
||||
|
||||
it('keeps the model that was already selected', async () => {
|
||||
mockApi.getModels.mockResolvedValue(models);
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
|
||||
await act(async () => {
|
||||
await result.current.fetchModels();
|
||||
});
|
||||
expect(state().activeModelId).toBe('model-a');
|
||||
});
|
||||
|
||||
it('does not select anything when the server has no models', async () => {
|
||||
useAppStore.setState({ activeModelId: '' });
|
||||
mockApi.getModels.mockResolvedValue([]);
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
|
||||
await act(async () => {
|
||||
await result.current.fetchModels();
|
||||
});
|
||||
expect(state().activeModelId).toBe('');
|
||||
});
|
||||
|
||||
it('marks the server offline and clears the list when the request fails', async () => {
|
||||
useAppStore.setState({ isConnected: true, availableModels: models });
|
||||
mockApi.getModels.mockRejectedValue(new Error('down'));
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
|
||||
let returned: unknown;
|
||||
await act(async () => {
|
||||
returned = await result.current.fetchModels();
|
||||
});
|
||||
expect(returned).toEqual([]);
|
||||
expect(state().isConnected).toBe(false);
|
||||
expect(state().availableModels).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('sendMessage', () => {
|
||||
it('streams an answer for the conversation, prefixed by the system prompt', async () => {
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
await act(async () => {
|
||||
await result.current.sendMessage('hello');
|
||||
});
|
||||
|
||||
const [, modelId, apiMessages, onChunk, onDone] = lastStreamCall();
|
||||
expect(modelId).toBe('model-a');
|
||||
expect(apiMessages[0].role).toBe('system');
|
||||
expect(apiMessages.slice(1)).toEqual([{ role: 'user', content: 'hello' }]);
|
||||
expect(state().isStreaming).toBe(true);
|
||||
|
||||
act(() => {
|
||||
onChunk('Hi');
|
||||
onChunk(' there');
|
||||
onDone();
|
||||
});
|
||||
expect(state().messages.map((m) => `${m.role}:${m.content}`)).toEqual(['user:hello', 'assistant:Hi there']);
|
||||
expect(state().isStreaming).toBe(false);
|
||||
expect(state().chatSessions).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('asks for a model instead of calling the server when none is selected', async () => {
|
||||
useAppStore.setState({ activeModelId: '' });
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
await act(async () => {
|
||||
await result.current.sendMessage('hello');
|
||||
});
|
||||
|
||||
expect(mockApi.chatCompletionStream).not.toHaveBeenCalled();
|
||||
expect(state().messages[1].content).toContain('No model selected');
|
||||
expect(state().isStreaming).toBe(false);
|
||||
});
|
||||
|
||||
it('ignores new messages while an answer is streaming', async () => {
|
||||
useAppStore.setState({ isStreaming: true });
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
await act(async () => {
|
||||
await result.current.sendMessage('again');
|
||||
});
|
||||
|
||||
expect(state().messages).toEqual([]);
|
||||
expect(mockApi.chatCompletionStream).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('writes the error into the assistant message and stops streaming when the stream fails', async () => {
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
await act(async () => {
|
||||
await result.current.sendMessage('hello');
|
||||
});
|
||||
|
||||
const onError = lastStreamCall()[5];
|
||||
act(() => {
|
||||
onError(new Error('Could not reach the server'));
|
||||
});
|
||||
expect(state().messages[1].content).toContain('Error: Could not reach the server');
|
||||
expect(state().isStreaming).toBe(false);
|
||||
expect(state().chatSessions).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('reports a synchronous failure to start the stream', async () => {
|
||||
mockApi.chatCompletionStream.mockImplementation(() => {
|
||||
throw new Error('boom');
|
||||
});
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
await act(async () => {
|
||||
await result.current.sendMessage('hello');
|
||||
});
|
||||
|
||||
expect(state().messages[1].content).toContain('Connection failed: boom');
|
||||
expect(state().isStreaming).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('regenerateResponse', () => {
|
||||
it('drops everything after the last user message and asks again', async () => {
|
||||
state().addMessage({ role: 'user', content: 'q1' });
|
||||
state().addMessage({ role: 'assistant', content: 'a1' });
|
||||
state().addMessage({ role: 'user', content: 'q2' });
|
||||
state().addMessage({ role: 'assistant', content: 'a2' });
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
|
||||
await act(async () => {
|
||||
await result.current.regenerateResponse();
|
||||
});
|
||||
const apiMessages = lastStreamCall()[2];
|
||||
expect(apiMessages.slice(1).map((m) => m.content)).toEqual(['q1', 'a1', 'q2']);
|
||||
expect(state().messages.map((m) => m.content)).toEqual(['q1', 'a1', 'q2', '']);
|
||||
expect(state().isStreaming).toBe(true);
|
||||
});
|
||||
|
||||
it('does nothing without a user message, while streaming, or without a model', async () => {
|
||||
const { result, rerender } = renderHook(() => useLMStudio());
|
||||
await act(async () => {
|
||||
await result.current.regenerateResponse();
|
||||
});
|
||||
expect(mockApi.chatCompletionStream).not.toHaveBeenCalled();
|
||||
|
||||
act(() => {
|
||||
state().addMessage({ role: 'user', content: 'q' });
|
||||
useAppStore.setState({ activeModelId: '' });
|
||||
});
|
||||
await act(async () => {
|
||||
await result.current.regenerateResponse();
|
||||
});
|
||||
expect(mockApi.chatCompletionStream).not.toHaveBeenCalled();
|
||||
|
||||
act(() => {
|
||||
useAppStore.setState({ activeModelId: 'model-a', isStreaming: true });
|
||||
});
|
||||
rerender({});
|
||||
await act(async () => {
|
||||
await result.current.regenerateResponse();
|
||||
});
|
||||
expect(mockApi.chatCompletionStream).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('stopStreaming', () => {
|
||||
it('aborts the request, stops streaming and keeps the partial answer in History', async () => {
|
||||
const abort = jest.fn();
|
||||
mockApi.chatCompletionStream.mockReturnValue({ abort });
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
await act(async () => {
|
||||
await result.current.sendMessage('hello');
|
||||
});
|
||||
act(() => {
|
||||
lastStreamCall()[3]('partial');
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stopStreaming();
|
||||
});
|
||||
expect(abort).toHaveBeenCalledTimes(1);
|
||||
expect(state().isStreaming).toBe(false);
|
||||
expect(state().chatSessions[0].messages.at(-1)?.content).toBe('partial');
|
||||
});
|
||||
|
||||
it('is a no-op when nothing is streaming', () => {
|
||||
const { result } = renderHook(() => useLMStudio());
|
||||
act(() => {
|
||||
result.current.stopStreaming();
|
||||
});
|
||||
expect(state().chatSessions).toEqual([]);
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user