refactor(frontend): Refactor tests to use Zustand's native state setting instead of vi.mock for stores (#11980)

Co-authored-by: amanape <83104063+amanape@users.noreply.github.com>

frontend/__tests__/components/browser.test.tsx

@@ -30,61 +30,33 @@ vi.mock("react-i18next", async () => {
   };
 });
 
-// Mock Zustand browser store
-let mockBrowserState = {
-  url: "https://example.com",
-  screenshotSrc: "",
-  setUrl: vi.fn(),
-  setScreenshotSrc: vi.fn(),
-  reset: vi.fn(),
-};
-
-vi.mock("#/stores/browser-store", () => ({
-  useBrowserStore: () => mockBrowserState,
-}));
-
-// Import the component after all mocks are set up
 import { BrowserPanel } from "#/components/features/browser/browser";
+import { useBrowserStore } from "#/stores/browser-store";
 
 describe("Browser", () => {
   afterEach(() => {
     vi.clearAllMocks();
-    // Reset the mock state
-    mockBrowserState = {
-      url: "https://example.com",
-      screenshotSrc: "",
-      setUrl: vi.fn(),
-      setScreenshotSrc: vi.fn(),
-      reset: vi.fn(),
-    };
   });
 
   it("renders a message if no screenshotSrc is provided", () => {
-    // Set the mock state for this test
-    mockBrowserState = {
+    useBrowserStore.setState({
       url: "https://example.com",
       screenshotSrc: "",
-      setUrl: vi.fn(),
-      setScreenshotSrc: vi.fn(),
       reset: vi.fn(),
-    };
+    });
 
     render(<BrowserPanel />);
 
-    // i18n empty message key
     expect(screen.getByText("BROWSER$NO_PAGE_LOADED")).toBeInTheDocument();
   });
 
   it("renders the url and a screenshot", () => {
-    // Set the mock state for this test
-    mockBrowserState = {
+    useBrowserStore.setState({
       url: "https://example.com",
       screenshotSrc:
         "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mN0uGvyHwAFCAJS091fQwAAAABJRU5ErkJggg==",
-      setUrl: vi.fn(),
-      setScreenshotSrc: vi.fn(),
       reset: vi.fn(),
-    };
+    });
 
     render(<BrowserPanel />);
 

frontend/__tests__/components/chat/chat-interface.test.tsx

@@ -25,10 +25,7 @@ import { useUnifiedUploadFiles } from "#/hooks/mutation/use-unified-upload-files
 import { OpenHandsAction } from "#/types/core/actions";
 import { useEventStore } from "#/stores/use-event-store";
 
-// Mock the hooks
 vi.mock("#/context/ws-client-provider");
-vi.mock("#/stores/error-message-store");
-vi.mock("#/stores/optimistic-user-message-store");
 vi.mock("#/hooks/query/use-config");
 vi.mock("#/hooks/mutation/use-get-trajectory");
 vi.mock("#/hooks/mutation/use-unified-upload-files");
@@ -102,24 +99,20 @@ describe("ChatInterface - Chat Suggestions", () => {
       },
     });
 
-    // Default mock implementations
     (useWsClient as unknown as ReturnType<typeof vi.fn>).mockReturnValue({
       send: vi.fn(),
       isLoadingMessages: false,
       parsedEvents: [],
     });
-    (
-      useOptimisticUserMessageStore as unknown as ReturnType<typeof vi.fn>
-    ).mockReturnValue({
-      setOptimisticUserMessage: vi.fn(),
-      getOptimisticUserMessage: vi.fn(() => null),
+
+    useOptimisticUserMessageStore.setState({
+      optimisticUserMessage: null,
     });
-    (
-      useErrorMessageStore as unknown as ReturnType<typeof vi.fn>
-    ).mockReturnValue({
-      setErrorMessage: vi.fn(),
-      removeErrorMessage: vi.fn(),
+
+    useErrorMessageStore.setState({
+      errorMessage: null,
     });
+
     (useConfig as unknown as ReturnType<typeof vi.fn>).mockReturnValue({
       data: { APP_MODE: "local" },
     });
@@ -204,11 +197,8 @@ describe("ChatInterface - Chat Suggestions", () => {
   });
 
   test("should hide chat suggestions when there is an optimistic user message", () => {
-    (
-      useOptimisticUserMessageStore as unknown as ReturnType<typeof vi.fn>
-    ).mockReturnValue({
-      setOptimisticUserMessage: vi.fn(),
-      getOptimisticUserMessage: vi.fn(() => "Optimistic message"),
+    useOptimisticUserMessageStore.setState({
+      optimisticUserMessage: "Optimistic message",
     });
 
     renderWithQueryClient(<ChatInterface />, queryClient);
@@ -240,24 +230,19 @@ describe("ChatInterface - Empty state", () => {
   });
 
   beforeEach(() => {
-    // Reset mocks to ensure empty state
     (useWsClient as unknown as ReturnType<typeof vi.fn>).mockReturnValue({
       send: sendMock,
       status: "CONNECTED",
       isLoadingMessages: false,
       parsedEvents: [],
     });
-    (
-      useOptimisticUserMessageStore as unknown as ReturnType<typeof vi.fn>
-    ).mockReturnValue({
-      setOptimisticUserMessage: vi.fn(),
-      getOptimisticUserMessage: vi.fn(() => null),
+
+    useOptimisticUserMessageStore.setState({
+      optimisticUserMessage: null,
     });
-    (
-      useErrorMessageStore as unknown as ReturnType<typeof vi.fn>
-    ).mockReturnValue({
-      setErrorMessage: vi.fn(),
-      removeErrorMessage: vi.fn(),
+
+    useErrorMessageStore.setState({
+      errorMessage: null,
     });
     (useConfig as unknown as ReturnType<typeof vi.fn>).mockReturnValue({
       data: { APP_MODE: "local" },

frontend/__tests__/components/interactive-chat-box.test.tsx

@@ -8,16 +8,10 @@ import { AgentState } from "#/types/agent-state";
 import { useAgentState } from "#/hooks/use-agent-state";
 import { useConversationStore } from "#/state/conversation-store";
 
-// Mock the agent state hook
 vi.mock("#/hooks/use-agent-state", () => ({
   useAgentState: vi.fn(),
 }));
 
-// Mock the conversation store
-vi.mock("#/state/conversation-store", () => ({
-  useConversationStore: vi.fn(),
-}));
-
 // Mock React Router hooks
 vi.mock("react-router", async () => {
   const actual = await vi.importActual("react-router");
@@ -58,44 +52,23 @@ vi.mock("#/hooks/use-conversation-name-context-menu", () => ({
 describe("InteractiveChatBox", () => {
   const onSubmitMock = vi.fn();
 
-  // Helper function to mock stores
   const mockStores = (agentState: AgentState = AgentState.INIT) => {
     vi.mocked(useAgentState).mockReturnValue({
       curAgentState: agentState,
     });
 
-    vi.mocked(useConversationStore).mockReturnValue({
+    useConversationStore.setState({
       images: [],
       files: [],
-      addImages: vi.fn(),
-      addFiles: vi.fn(),
-      clearAllFiles: vi.fn(),
-      addFileLoading: vi.fn(),
-      removeFileLoading: vi.fn(),
-      addImageLoading: vi.fn(),
-      removeImageLoading: vi.fn(),
-      submittedMessage: null,
-      setShouldHideSuggestions: vi.fn(),
-      setSubmittedMessage: vi.fn(),
-      isRightPanelShown: true,
-      selectedTab: "editor" as const,
       loadingFiles: [],
       loadingImages: [],
+      submittedMessage: null,
       messageToSend: null,
       shouldShownAgentLoading: false,
       shouldHideSuggestions: false,
+      isRightPanelShown: true,
+      selectedTab: "editor" as const,
       hasRightPanelToggled: true,
-      setIsRightPanelShown: vi.fn(),
-      setSelectedTab: vi.fn(),
-      setShouldShownAgentLoading: vi.fn(),
-      removeImage: vi.fn(),
-      removeFile: vi.fn(),
-      clearImages: vi.fn(),
-      clearFiles: vi.fn(),
-      clearAllLoading: vi.fn(),
-      setMessageToSend: vi.fn(),
-      resetConversationState: vi.fn(),
-      setHasRightPanelToggled: vi.fn(),
     });
   };
 

frontend/__tests__/services/actions.test.tsx

@@ -1,8 +1,8 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import ActionType from "#/types/action-type";
 import { ActionMessage } from "#/types/message";
+import { useCommandStore } from "#/state/command-store";
 
-// Mock the store and actions
 const mockDispatch = vi.fn();
 const mockAppendInput = vi.fn();
 
@@ -12,26 +12,12 @@ vi.mock("#/store", () => ({
   },
 }));
 
-vi.mock("#/state/command-store", () => ({
-  useCommandStore: {
-    getState: () => ({
-      appendInput: mockAppendInput,
-    }),
-  },
-}));
-
-vi.mock("#/state/metrics-slice", () => ({
-  setMetrics: vi.fn(),
-}));
-
-vi.mock("#/state/security-analyzer-slice", () => ({
-  appendSecurityAnalyzerInput: vi.fn(),
-}));
-
 describe("handleActionMessage", () => {
   beforeEach(() => {
-    // Clear all mocks before each test
     vi.clearAllMocks();
+    useCommandStore.setState({
+      appendInput: mockAppendInput,
+    });
   });
 
   it("should handle RUN actions by adding input to terminal", async () => {