Instruction file imported from vual/lobe-chat-pro (.cursor/rules/testing-guide/zustand-store-action-test.mdc). Copyright stays with the author.

Zustand Store Action Testing Guide

This guide provides best practices for testing Zustand store actions, based on our proven testing patterns.

Basic Test Structure

import { act, renderHook } from '@testing-library/react';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

import { messageService } from '@/services/message';
import { useChatStore } from '../../store';

// Keep zustand mock as it's needed globally
vi.mock('zustand/traditional');

beforeEach(() => {
  // Reset store state
  vi.clearAllMocks();
  useChatStore.setState(
    {
      activeId: 'test-session-id',
      messagesMap: {},
      loadingIds: [],
    },
    false,
  );

  // ✅ Setup only spies that MOST tests need
  vi.spyOn(messageService, 'createMessage').mockResolvedValue('new-message-id');
  // ❌ Don't setup spies that only few tests need - spy only when needed

  // Setup common mock methods
  act(() => {
    useChatStore.setState({
      refreshMessages: vi.fn(),
      internal_coreProcessMessage: vi.fn(),
    });
  });
});

afterEach(() => {
  vi.restoreAllMocks();
});

describe('action name', () => {
  describe('validation', () => {
    // Validation tests
  });

  describe('normal flow', () => {
    // Happy path tests
  });

  describe('error handling', () => {
    // Error case tests
  });
});

Testing Best Practices

1. Test Layering - Spy Direct Dependencies Only

✅ Good: Spy on the direct dependency

// When testing internal_coreProcessMessage, spy its direct dependency
const fetchAIChatSpy = vi
  .spyOn(result.current, 'internal_fetchAIChatMessage')
  .mockResolvedValue({ isFunctionCall: false, content: 'AI response' });

❌ Bad: Spy on lower-level implementation details

// Don't spy on services that internal_fetchAIChatMessage uses
const streamSpy = vi
  .spyOn(chatService, 'createAssistantMessageStream')
  .mockImplementation(...);

Why: Each test should only mock its direct dependencies, not the entire call chain. This makes tests more maintainable and less brittle.

2. Mock Management - Minimize Global Spies

✅ Good: Spy only when needed

beforeEach(() => {
  // ✅ Only spy services that most tests need
  vi.spyOn(messageService, 'createMessage').mockResolvedValue('new-message-id');
  // ✅ Don't spy chatService globally
});

it('should process message', async () => {
  // ✅ Spy chatService only in tests that need it
  const streamSpy = vi.spyOn(chatService, 'createAssistantMessageStream')
    .mockImplementation(...);

  // test logic

  streamSpy.mockRestore();
});

❌ Bad: Setup all spies globally

beforeEach(() => {
  vi.spyOn(messageService, 'createMessage').mockResolvedValue('id');
  vi.spyOn(chatService, 'createAssistantMessageStream').mockResolvedValue({}); // ❌ Not all tests need this
  vi.spyOn(fileService, 'uploadFile').mockResolvedValue({}); // ❌ Creates implicit coupling
});

3. Service Mocking - Mock the Correct Layer

✅ Good: Mock the service method

it('should fetch AI chat response', async () => {
  const streamSpy = vi
    .spyOn(chatService, 'createAssistantMessageStream')
    .mockImplementation(async ({ onMessageHandle, onFinish }) => {
      await onMessageHandle?.({ type: 'text', text: 'Hello' } as any);
      await onFinish?.('Hello', {});
    });

  // test logic
});

❌ Bad: Mock global fetch

it('should fetch AI chat response', async () => {
  global.fetch = vi.fn().mockResolvedValue(...); // ❌ Too low level
});

4. Test Organization - Use Descriptive Nesting

✅ Good: Clear nested structure

describe('sendMessage', () => {
  describe('validation', () => {
    it('should not send when session is inactive', async () => {});
    it('should not send when message is empty', async () => {});
  });

  describe('message creation', () => {
    it('should create user message and trigger AI processing', async () => {});
    it('should send message with files attached', async () => {});
  });

  describe('error handling', () => {
    it('should handle message creation errors gracefully', async () => {});
  });
});

❌ Bad: Flat structure

describe('sendMessage', () => {
  it('test 1', async () => {});
  it('test 2', async () => {});
  it('test 3', async () => {});
});

5. Testing Async Actions

Always wrap async operations in act():

it('should send message', async () => {
  const { result } = renderHook(() => useChatStore());

  await act(async () => {
    await result.current.sendMessage({ message: 'Hello' });
  });

  expect(messageService.createMessage).toHaveBeenCalled();
});

6. State Setup - Use act() for setState

it('should handle disabled state', async () => {
  act(() => {
    useChatStore.setState({ activeId: undefined });
  });

  const { result } = renderHook(() => useChatStore());
  // test logic
});

7. Testing Complex Flows

For complex flows with multiple steps, use clear spy setup:

it('should handle topic creation flow', async () => {
  // Setup store state
  act(() => {
    useChatStore.setState({
      activeTopicId: undefined,
      messagesMap: {
        'test-session-id': [
          { id: 'msg-1', role: 'user', content: 'Message 1' },
          { id: 'msg-2', role: 'assistant', content: 'Response 1' },
          { id: 'msg-3', role: 'user', content: 'Message 2' },
        ],
      },
    });
  });

  const { result } = renderHook(() => useChatStore());

  // Spy on action dependencies
  const createTopicSpy = vi.spyOn(result.current, 'createTopic')
    .mockResolvedValue('new-topic-id');
  const toggleLoadingSpy = vi.spyOn(result.current, 'internal_toggleMessageLoading');

  // Execute
  await act(async () => {
    await result.current.sendMessage({ message: 'Test message' });
  });

  // Assert
  expect(createTopicSpy).toHaveBeenCalled();
  expect(toggleLoadingSpy).toHaveBeenCalledWith(true, expect.any(String));
});

8. Streaming Response Mocking

When testing streaming responses, simulate the flow properly:

it('should handle streaming chunks', async () => {
  const { result } = renderHook(() => useChatStore());
  const messages = [
    { id: 'msg-1', role: 'user', content: 'Hello', sessionId: 'test-session' },
  ];

  const streamSpy = vi
    .spyOn(chatService, 'createAssistantMessageStream')
    .mockImplementation(async ({ onMessageHandle, onFinish }) => {
      // Simulate streaming chunks
      await onMessageHandle?.({ type: 'text', text: 'Hello' } as any);
      await onMessageHandle?.({ type: 'text', text: ' World' } as any);
      await onFinish?.('Hello World', {});
    });

  await act(async () => {
    await result.current.internal_fetchAIChatMessage({
      messages,
      messageId: 'test-message-id',
      model: 'gpt-4o-mini',
      provider: 'openai',
    });
  });

  expect(result.current.internal_dispatchMessage).toHaveBeenCalled();

  streamSpy.mockRestore();
});

9. Error Handling Tests

Always test error scenarios:

it('should handle errors gracefully', async () => {
  const { result } = renderHook(() => useChatStore());

  vi.spyOn(messageService, 'createMessage').mockRejectedValue(
    new Error('create message error'),
  );

  await act(async () => {
    try {
      await result.current.sendMessage({ message: 'Test message' });
    } catch {
      // Expected to throw
    }
  });

  expect(result.current.internal_coreProcessMessage).not.toHaveBeenCalled();
});

10. Cleanup After Tests

Always restore mocks after each test:

afterEach(() => {
  vi.restoreAllMocks();
});

// For individual test cleanup:
it('should test something', async () => {
  const spy = vi.spyOn(service, 'method').mockImplementation(...);

  // test logic

  spy.mockRestore(); // Optional: cleanup immediately after test
});

Common Patterns

Testing Store Methods That Call Other Store Methods

it('should call internal methods', async () => {
  const { result } = renderHook(() => useChatStore());

  const internalMethodSpy = vi.spyOn(result.current, 'internal_method')
    .mockResolvedValue();

  await act(async () => {
    await result.current.publicMethod();
  });

  expect(internalMethodSpy).toHaveBeenCalledWith(
    expect.any(String),
    expect.objectContaining({ key: 'value' }),
  );
});

Testing Conditional Logic

describe('conditional behavior', () => {
  it('should execute when condition is true', async () => {
    const { result } = renderHook(() => useChatStore());
    vi.spyOn(result.current, 'internal_shouldUseRAG').mockReturnValue(true);

    await act(async () => {
      await result.current.sendMessage({ message: 'test' });
    });

    expect(result.current.internal_retrieveChunks).toHaveBeenCalled();
  });

  it('should not execute when condition is false', async () => {
    const { result } = renderHook(() => useChatStore());
    vi.spyOn(result.current, 'internal_shouldUseRAG').mockReturnValue(false);

    await act(async () => {
      await result.current.sendMessage({ message: 'test' });
    });

    expect(result.current.internal_retrieveChunks).not.toHaveBeenCalled();
  });
});

Testing AbortController

it('should abort generation and clear loading state', () => {
  const abortController = new AbortController();

  act(() => {
    useChatStore.setState({ chatLoadingIdsAbortController: abortController });
  });

  const { result } = renderHook(() => useChatStore());
  const toggleLoadingSpy = vi.spyOn(result.current, 'internal_toggleChatLoading');

  act(() => {
    result.current.stopGenerateMessage();
  });

  expect(abortController.signal.aborted).toBe(true);
  expect(toggleLoadingSpy).toHaveBeenCalledWith(false, undefined, expect.any(String));
});

Anti-Patterns to Avoid

❌ Don't: Mock the entire store

vi.mock('../../store', () => ({
  useChatStore: vi.fn(() => ({
    sendMessage: vi.fn(),
  })),
}));

❌ Don't: Test implementation details

// Bad: testing internal state structure
expect(result.current.messagesMap).toHaveProperty('test-session');

// Good: testing behavior
expect(result.current.refreshMessages).toHaveBeenCalled();

❌ Don't: Create tight coupling between tests

// Bad: Tests depend on order
let messageId: string;

it('test 1', () => {
  messageId = 'some-id'; // Side effect
});

it('test 2', () => {
  expect(messageId).toBeDefined(); // Depends on test 1
});

❌ Don't: Over-mock services

// Bad: Mocking everything
beforeEach(() => {
  vi.mock('@/services/chat');
  vi.mock('@/services/message');
  vi.mock('@/services/file');
  vi.mock('@/services/agent');
  // ... too many global mocks
});

Testing SWR Hooks in Zustand Stores

Some Zustand store slices use SWR hooks for data fetching. These require a different testing approach.

Basic SWR Hook Test Structure

import { renderHook, waitFor } from '@testing-library/react';
import { beforeEach, describe, expect, it, vi } from 'vitest';

import { discoverService } from '@/services/discover';
import { globalHelpers } from '@/store/global/helpers';
import { useDiscoverStore as useStore } from '../../store';

vi.mock('zustand/traditional');

beforeEach(() => {
  vi.clearAllMocks();
});

describe('SWR Hook Actions', () => {
  it('should fetch data and return correct response', async () => {
    const mockData = [{ id: '1', name: 'Item 1' }];

    // Mock the service call (the fetcher)
    vi.spyOn(discoverService, 'getPluginCategories').mockResolvedValue(mockData as any);
    vi.spyOn(globalHelpers, 'getCurrentLanguage').mockReturnValue('en-US');

    const params = {} as any;
    const { result } = renderHook(() => useStore.getState().usePluginCategories(params));

    // Use waitFor to wait for async data loading
    await waitFor(() => {
      expect(result.current.data).toEqual(mockData);
    });

    expect(discoverService.getPluginCategories).toHaveBeenCalledWith(params);
  });
});

Key points:

• DO NOT mock useSWR - let it use the real implementation
• Only mock the service methods (fetchers)
• Use waitFor from @testing-library/react to wait for async operations
• Check result.current.data directly after waitFor completes

Testing SWR Key Generation

it('should generate correct SWR key with locale and params', () => {
  vi.spyOn(globalHelpers, 'getCurrentLanguage').mockReturnValue('zh-CN');

  const useSWRMock = vi.mocked(useSWR);
  let capturedKey: string | null = null;
  useSWRMock.mockImplementation(((key: string) => {
    capturedKey = key;
    return { data: undefined, error: undefined, isValidating: false, mutate: vi.fn() };
  }) as any);

  const params = { page: 2, category: 'tools' } as any;
  renderHook(() => useStore.getState().usePluginList(params));

  expect(capturedKey).toBe('plugin-list-zh-CN-2-tools');
});

Testing SWR Configuration

it('should have correct SWR configuration', () => {
  const useSWRMock = vi.mocked(useSWR);
  let capturedOptions: any = null;
  useSWRMock.mockImplementation(((key: string, fetcher: any, options: any) => {
    capturedOptions = options;
    return { data: undefined, error: undefined, isValidating: false, mutate: vi.fn() };
  }) as any);

  renderHook(() => useStore.getState().usePluginIdentifiers());

  expect(capturedOptions).toMatchObject({ revalidateOnFocus: false });
});

Testing Conditional Fetching

it('should not fetch when required parameter is missing', () => {
  const useSWRMock = vi.mocked(useSWR);
  let capturedKey: string | null = null;
  useSWRMock.mockImplementation(((key: string | null) => {
    capturedKey = key;
    return { data: undefined, error: undefined, isValidating: false, mutate: vi.fn() };
  }) as any);

  // When identifier is undefined, SWR key should be null
  renderHook(() => useStore.getState().usePluginDetail({ identifier: undefined }));

  expect(capturedKey).toBeNull();
});

Key Differences from Regular Action Tests

1. Mock useSWR globally: Use vi.mock('swr') at the top level
2. Mock the fetcher, not the result:
   • ✅ Correct: const data = fetcher?.() - call fetcher and return its Promise
   • ❌ Wrong: return { data: mockData } - hardcode the result
3. Await Promise results: The data field is a Promise, use await result.current.data
4. No act() wrapper needed: SWR hooks don't trigger React state updates in these tests
5. Test SWR key generation: Verify keys include locale and parameters
6. Test configuration: Verify revalidation and other SWR options
7. Type assertions: Use as any for test mock data where type definitions are strict

Why this matters:

• The fetcher (service method) is what we're testing - it must be called
• Hardcoding the return value bypasses the actual fetcher logic
• SWR returns Promises in real usage, tests should mirror this behavior

Benefits of This Approach

✅ Clear test layers - Each test only spies on direct dependencies ✅ Correct mocks - Mocks match actual implementation ✅ Better maintainability - Changes to implementation require fewer test updates ✅ Improved coverage - Structured approach ensures all branches are tested ✅ Reduced coupling - Tests are independent and can run in any order

Reference

See example implementation in:

• src/store/chat/slices/aiChat/actions/__tests__/generateAIChat.test.ts (Regular actions)
• src/store/discover/slices/plugin/action.test.ts (SWR hooks)
• src/store/discover/slices/mcp/action.test.ts (SWR hooks)