LocalMind/__tests__/useLMStudio.test.ts

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('keeps the current answer and asks for a model when none is selected', async () => {
        state().addMessage({ role: 'user', content: 'q' });
        state().addMessage({ role: 'assistant', content: 'a' });
        useAppStore.setState({ activeModelId: '' });
        const { result } = renderHook(() => useLMStudio());

        await act(async () => {
            await result.current.regenerateResponse();
        });
        expect(mockApi.chatCompletionStream).not.toHaveBeenCalled();
        expect(state().messages.map((m) => m.content)).toEqual([
            'q',
            'a',
            expect.stringContaining('No model selected'),
        ]);
    });

    it('does nothing without a user message or while streaming', async () => {
        const { result, rerender } = renderHook(() => useLMStudio());
        await act(async () => {
            await result.current.regenerateResponse();
        });
        expect(mockApi.chatCompletionStream).not.toHaveBeenCalled();
        expect(state().messages).toEqual([]);

        act(() => {
            state().addMessage({ role: 'user', content: 'q' });
        });
        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([]);
    });
});