fix: updated broken tests

src/components/Dialog/DeleteMessagesDialog/DeleteMessagesDialog.test.tsx

@@ -1,9 +1,10 @@
 import { render, screen, fireEvent } from '@testing-library/react';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
+// Ensure the path is correct for import
 import { useMessageStore } from "../../../core/stores/messageStore/index.ts";
 import { DeleteMessagesDialog } from "@components/Dialog/DeleteMessagesDialog/DeleteMessagesDialog.tsx";
 
-vi.mock('@core/stores/messageStore.ts', () => ({
+vi.mock('@core/stores/messageStore', () => ({
   useMessageStore: vi.fn(() => ({
     deleteAllMessages: vi.fn(),
   })),
@@ -14,17 +15,36 @@ describe('DeleteMessagesDialog', () => {
   const mockClearAllMessages = vi.fn();
 
   beforeEach(() => {
-    vi.mocked(useMessageStore).mockReturnValue({ deleteAllMessages: mockClearAllMessages });
     mockOnOpenChange.mockClear();
     mockClearAllMessages.mockClear();
+
+    const mockedUseMessageStore = vi.mocked(useMessageStore);
+    mockedUseMessageStore.mockImplementation(() => ({
+      deleteAllMessages: mockClearAllMessages
+    }));
+    mockedUseMessageStore.mockClear();
+
+  });
+
+  it('calls onOpenChange with false when the close button (X) is clicked', () => {
+    render(<DeleteMessagesDialog open={true} onOpenChange={mockOnOpenChange} />);
+    const closeButton = screen.queryByTestId('dialog-close-button');
+    if (!closeButton) {
+      throw new Error("Dialog close button with data-testid='dialog-close-button' not found. Did you add it to the component?");
+    }
+    fireEvent.click(closeButton);
+    expect(mockOnOpenChange).toHaveBeenCalledTimes(1);
+    expect(mockOnOpenChange).toHaveBeenCalledWith(false);
   });
 
+
+
   it('renders the dialog when open is true', () => {
     render(<DeleteMessagesDialog open={true} onOpenChange={mockOnOpenChange} />);
-    expect(screen.getByText('Clear All Messages')).toBeVisible();
-    expect(screen.getByText(/This action will clear all message history./)).toBeVisible();
-    expect(screen.getByRole('button', { name: 'Dismiss' })).toBeVisible();
-    expect(screen.getByRole('button', { name: 'Clear Messages' })).toBeVisible();
+    expect(screen.getByText('Clear All Messages')).toBeInTheDocument();
+    expect(screen.getByText(/This action will clear all message history./)).toBeInTheDocument();
+    expect(screen.getByRole('button', { name: 'Dismiss' })).toBeInTheDocument();
+    expect(screen.getByRole('button', { name: 'Clear Messages' })).toBeInTheDocument();
   });
 
   it('does not render the dialog when open is false', () => {
@@ -32,15 +52,10 @@ describe('DeleteMessagesDialog', () => {
     expect(screen.queryByText('Clear All Messages')).toBeNull();
   });
 
-  it('calls onOpenChange with false when the close button is clicked', () => {
-    render(<DeleteMessagesDialog open={true} onOpenChange={mockOnOpenChange} />);
-    fireEvent.click(screen.getByRole('button', { name: 'Close' }));
-    expect(mockOnOpenChange).toHaveBeenCalledWith(false);
-  });
-
   it('calls onOpenChange with false when the dismiss button is clicked', () => {
     render(<DeleteMessagesDialog open={true} onOpenChange={mockOnOpenChange} />);
     fireEvent.click(screen.getByRole('button', { name: 'Dismiss' }));
+    expect(mockOnOpenChange).toHaveBeenCalledTimes(1); // Add count check
     expect(mockOnOpenChange).toHaveBeenCalledWith(false);
   });
 
@@ -48,6 +63,7 @@ describe('DeleteMessagesDialog', () => {
     render(<DeleteMessagesDialog open={true} onOpenChange={mockOnOpenChange} />);
     fireEvent.click(screen.getByRole('button', { name: 'Clear Messages' }));
     expect(mockClearAllMessages).toHaveBeenCalledTimes(1);
+    expect(mockOnOpenChange).toHaveBeenCalledTimes(1); // Add count check
     expect(mockOnOpenChange).toHaveBeenCalledWith(false);
   });
 });

src/components/Dialog/DeleteMessagesDialog/DeleteMessagesDialog.tsx

@@ -29,7 +29,7 @@ export const DeleteMessagesDialog = ({
   return (
     <Dialog open={open} onOpenChange={onOpenChange}>
       <DialogContent>
-        <DialogClose onClick={handleCloseDialog} />
+        <DialogClose data-testid="dialog-close-button" />
         <DialogHeader>
           <DialogTitle className="flex items-center gap-2">
             <AlertTriangleIcon className="h-5 w-5 text-warning" />

src/components/Dialog/RefreshKeysDialog/useRefreshKeysDialog.test.ts

@@ -2,12 +2,12 @@ import { renderHook, act } from "@testing-library/react";
 import { useRefreshKeysDialog } from "./useRefreshKeysDialog.ts";
 import { beforeEach, describe, expect, it, Mock, vi } from "vitest";
 import { useDevice } from "@core/stores/deviceStore.ts";
+import { useMessageStore } from "@core/stores/messageStore/index.ts";
 
-vi.mock("@core/stores/messageStore.ts", () => ({
+vi.mock("@core/stores/messageStore", () => ({
   useMessageStore: vi.fn(() => ({ activeChat: "chat-123" })),
 }));
-
-vi.mock("@core/stores/deviceStore.ts", () => ({
+vi.mock("@core/stores/deviceStore", () => ({
   useDevice: vi.fn(() => ({
     removeNode: vi.fn(),
     setDialogOpen: vi.fn(),
@@ -23,46 +23,50 @@ describe("useRefreshKeysDialog Hook", () => {
   let clearNodeErrorMock: Mock;
 
   beforeEach(() => {
+    vi.clearAllMocks();
+
     removeNodeMock = vi.fn();
     setDialogOpenMock = vi.fn();
-    getNodeErrorMock = vi.fn();
+    getNodeErrorMock = vi.fn().mockReturnValue(undefined);
     clearNodeErrorMock = vi.fn();
 
-    (useDevice as Mock).mockReturnValue({
+    vi.mocked(useDevice).mockReturnValue({
       removeNode: removeNodeMock,
       setDialogOpen: setDialogOpenMock,
       getNodeError: getNodeErrorMock,
       clearNodeError: clearNodeErrorMock,
     });
+
+    vi.mocked(useMessageStore).mockReturnValue({
+      activeChat: "chat-123"
+    });
   });
 
   it("handleNodeRemove should remove the node and update dialog if there is an error", () => {
     getNodeErrorMock.mockReturnValue({ node: "node-abc" });
 
     const { result } = renderHook(() => useRefreshKeysDialog());
+    act(() => { result.current.handleNodeRemove(); });
 
-    act(() => {
-      result.current.handleNodeRemove();
-    });
-
+    expect(getNodeErrorMock).toHaveBeenCalledTimes(1);
     expect(getNodeErrorMock).toHaveBeenCalledWith("chat-123");
+    expect(clearNodeErrorMock).toHaveBeenCalledTimes(1);
     expect(clearNodeErrorMock).toHaveBeenCalledWith("chat-123");
+    expect(removeNodeMock).toHaveBeenCalledTimes(1);
     expect(removeNodeMock).toHaveBeenCalledWith("node-abc");
+    expect(setDialogOpenMock).toHaveBeenCalledTimes(1);
     expect(setDialogOpenMock).toHaveBeenCalledWith("refreshKeys", false);
   });
 
   it("handleNodeRemove should do nothing if there is no error", () => {
-    getNodeErrorMock.mockReturnValue(undefined);
-
     const { result } = renderHook(() => useRefreshKeysDialog());
+    act(() => { result.current.handleNodeRemove(); });
 
-    act(() => {
-      result.current.handleNodeRemove();
-    });
-
+    expect(getNodeErrorMock).toHaveBeenCalledTimes(1);
+    expect(getNodeErrorMock).toHaveBeenCalledWith("chat-123");
+    expect(clearNodeErrorMock).not.toHaveBeenCalled();
     expect(removeNodeMock).not.toHaveBeenCalled();
     expect(setDialogOpenMock).not.toHaveBeenCalled();
-    expect(clearNodeErrorMock).not.toHaveBeenCalled();
   });
 
   it("handleCloseDialog should close the dialog", () => {
@@ -72,6 +76,7 @@ describe("useRefreshKeysDialog Hook", () => {
       result.current.handleCloseDialog();
     });
 
+    expect(setDialogOpenMock).toHaveBeenCalledTimes(1);
     expect(setDialogOpenMock).toHaveBeenCalledWith("refreshKeys", false);
   });
-});
+});

src/components/Dialog/RefreshKeysDialog/useRefreshKeysDialog.ts

@@ -1,28 +1,27 @@
 import { useCallback } from "react";
 import { useDevice } from "@core/stores/deviceStore.ts";
-import { useMessageStore } from "../../../core/stores/messageStore/index.ts";
+import { useMessageStore } from "@core/stores/messageStore/index.ts";
 
 export function useRefreshKeysDialog() {
   const { removeNode, setDialogOpen, clearNodeError, getNodeError } = useDevice();
   const { activeChat } = useMessageStore();
 
+  const handleCloseDialog = useCallback(() => {
+    setDialogOpen('refreshKeys', false);
+  }, [setDialogOpen]);
+
   const handleNodeRemove = useCallback(() => {
     const nodeWithError = getNodeError(activeChat);
     if (!nodeWithError) {
       return;
     }
     clearNodeError(activeChat);
-    handleCloseDialog();;
+    handleCloseDialog();
     return removeNode(nodeWithError?.node);
-  }, [activeChat, clearNodeError, setDialogOpen, removeNode]);
-
-  const handleCloseDialog = useCallback(() => {
-    setDialogOpen('refreshKeys', false);
-  }, [setDialogOpen])
+  }, [activeChat, clearNodeError, getNodeError, removeNode, handleCloseDialog]);
 
   return {
     handleCloseDialog,
     handleNodeRemove
   };
-
 }

src/components/PageComponents/Messages/MessageInput.test.tsx

@@ -1,14 +1,11 @@
-import { render, screen, fireEvent, waitFor } from '@testing-library/react';
+import { render, screen, fireEvent, waitFor, act } from '@testing-library/react';
 import { vi, describe, it, expect, beforeEach } from 'vitest';
-import { MessageInput } from './MessageInput.tsx';
-import { useDevice } from '@core/stores/deviceStore.ts';
-import { useMessageStore } from '../../../core/stores/messageStore/index.ts';
-import { debounce } from '@core/utils/debounce.ts';
-import { Types } from "@meshtastic/core";
+import { MessageInput, MessageInputProps } from './MessageInput.tsx';
+import { Types } from '@meshtastic/core';
 
 vi.mock('@components/UI/Button.tsx', () => ({
-  Button: vi.fn(({ type, className, children, onClick, onSubmit }) => (
-    <button type={type} className={className} onClick={onClick} onSubmit={onSubmit}>
+  Button: vi.fn(({ type, className, children, onClick, onSubmit, variant, ...rest }) => (
+    <button type={type} className={className} onClick={onClick} onSubmit={onSubmit} {...rest}>
       {children}
     </button>
   )),
@@ -23,16 +20,21 @@ vi.mock('@components/UI/Input.tsx', () => ({
       placeholder={placeholder}
       value={value}
       onChange={onChange}
+      data-testid="message-input-field"
     />
   )),
 }));
 
-vi.mock('@core/stores/deviceStore.ts', () => ({
-  useDevice: vi.fn(),
-}));
+const mockSetDraft = vi.fn();
+const mockGetDraft = vi.fn();
+const mockClearDraft = vi.fn();
 
-vi.mock('@core/stores/messageStore.ts', () => ({
-  useMessageStore: vi.fn(),
+vi.mock('@core/stores/messageStore', () => ({
+  useMessageStore: vi.fn(() => ({
+    setDraft: mockSetDraft,
+    getDraft: mockGetDraft,
+    clearDraft: mockClearDraft,
+  })),
   MessageState: {
     Ack: 'ack',
     Waiting: 'waiting',
@@ -44,111 +46,177 @@ vi.mock('@core/stores/messageStore.ts', () => ({
   },
 }));
 
-vi.mock('@core/utils/debounce.ts', () => ({
-  debounce: vi.fn((fn) => fn),
-}));
-
 vi.mock('lucide-react', () => ({
   SendIcon: vi.fn(() => <svg data-testid="send-icon" />),
 }));
 
 describe('MessageInput', () => {
-  const mockSetMessageState = vi.fn();
-  const mockSetActiveChat = vi.fn();
-  const mockSetDraft = vi.fn();
-  const mockGetDraft = vi.fn();
-  const mockClearDraft = vi.fn();
-  const mockSendText = vi.fn();
+  const mockOnSend = vi.fn();
+  const defaultProps: MessageInputProps = {
+    onSend: mockOnSend,
+    to: 123,
+    maxBytes: 256,
+  };
 
   beforeEach(() => {
-    (useDevice as ReturnType<typeof vi.fn>).mockReturnValue({
-      connection: {
-        sendText: mockSendText,
-      },
-    });
-
-    (useMessageStore as unknown as ReturnType<typeof vi.fn>).mockReturnValue({
-      setMessageState: mockSetMessageState,
-      activeChat: 123,
-      setDraft: mockSetDraft,
-      getDraft: mockGetDraft,
-      clearDraft: mockClearDraft,
-    });
+    vi.clearAllMocks();
 
-    mockSetMessageState.mockClear();
-    mockSetActiveChat.mockClear();
-    mockSetDraft.mockClear();
-    mockGetDraft.mockClear();
-    mockClearDraft.mockClear();
-    mockSendText.mockClear();
-    (debounce as ReturnType<typeof vi.fn>).mockImplementation((fn) => fn);
+    mockGetDraft.mockReturnValue('');
   });
 
-  const renderComponent = (props: { to: Types.Destination; channel: Types.ChannelNumber; maxBytes: number }) => {
-    render(<MessageInput {...props} />);
+  const renderComponent = (props: Partial<MessageInputProps> = {}) => {
+    render(<MessageInput {...defaultProps} {...props} />);
   };
 
-  it.skip('sends text message and updates state to Ack on submit', async () => {
-    renderComponent({ to: 2, channel: 3, maxBytes: 256 });
+  it('should render the input field, byte counter, and send button', () => {
+    renderComponent();
+    expect(screen.getByPlaceholderText('Enter Message')).toBeInTheDocument();
+    expect(screen.getByTestId('byte-counter')).toBeInTheDocument();
+    expect(screen.getByRole('button')).toBeInTheDocument();
+    expect(screen.getByTestId('send-icon')).toBeInTheDocument();
+  });
+
+  it('should initialize with the draft from the store', () => {
+    const initialDraft = 'Existing draft message';
+    mockGetDraft.mockImplementation((key) => {
+      return key === defaultProps.to ? initialDraft : '';
+    });
+
+    renderComponent();
+
     const inputElement = screen.getByPlaceholderText('Enter Message') as HTMLInputElement;
-    fireEvent.change(inputElement, { target: { value: 'Hello' } });
+    expect(inputElement.value).toBe(initialDraft);
+    expect(mockGetDraft).toHaveBeenCalledWith(defaultProps.to);
+    const expectedBytes = new Blob([initialDraft]).size;
+    expect(screen.getByTestId('byte-counter')).toHaveTextContent(`${expectedBytes}/${defaultProps.maxBytes}`);
+  });
+
+  it('should update input value, byte counter, and call setDraft on change within limits', () => {
+    renderComponent();
+    const inputElement = screen.getByPlaceholderText('Enter Message');
+    const testMessage = 'Hello there!';
+    const expectedBytes = new Blob([testMessage]).size;
+
+    fireEvent.change(inputElement, { target: { value: testMessage } });
+
+    expect((inputElement as HTMLInputElement).value).toBe(testMessage);
+    expect(screen.getByTestId('byte-counter')).toHaveTextContent(`${expectedBytes}/${defaultProps.maxBytes}`);
+    expect(mockSetDraft).toHaveBeenCalledTimes(1);
+    expect(mockSetDraft).toHaveBeenCalledWith(defaultProps.to, testMessage);
+  });
+
+  it('should NOT update input value or call setDraft if maxBytes is exceeded', () => {
+    const smallMaxBytes = 5;
+    renderComponent({ maxBytes: smallMaxBytes });
+    const inputElement = screen.getByPlaceholderText('Enter Message');
+    const initialValue = '12345';
+    const excessiveValue = '123456';
+
+    fireEvent.change(inputElement, { target: { value: initialValue } });
+    expect((inputElement as HTMLInputElement).value).toBe(initialValue);
+    expect(mockSetDraft).toHaveBeenCalledWith(defaultProps.to, initialValue);
+    mockSetDraft.mockClear();
+
+    fireEvent.change(inputElement, { target: { value: excessiveValue } });
+
+    expect((inputElement as HTMLInputElement).value).toBe(initialValue);
+    expect(screen.getByTestId('byte-counter')).toHaveTextContent(`${smallMaxBytes}/${smallMaxBytes}`);
+    expect(mockSetDraft).not.toHaveBeenCalled();
+  });
+
+  it('should call onSend, clear input, reset byte counter, and call clearDraft on valid submit', async () => {
+    renderComponent();
+    const inputElement = screen.getByPlaceholderText('Enter Message');
     const formElement = screen.getByRole('form');
+    const testMessage = 'Send this message';
+
+    fireEvent.change(inputElement, { target: { value: testMessage } });
     fireEvent.submit(formElement);
 
     await waitFor(() => {
-      expect(mockSendText).toHaveBeenCalledWith('Hello', 2, true, 3);
-      expect(mockSetMessageState).toHaveBeenCalledWith({
-        type: 'direct',
-        key: 123,
-        messageId: undefined,
-        newState: 'ack',
-      });
-      expect(mockClearDraft).toHaveBeenCalledWith(2);
-      expect(inputElement.value).toBe('');
-      expect(screen.getByTestId('byte-counter')).toHaveTextContent('0/256');
+      expect(mockOnSend).toHaveBeenCalledTimes(1);
+      expect(mockOnSend).toHaveBeenCalledWith(testMessage);
+      expect((inputElement as HTMLInputElement).value).toBe('');
+      expect(screen.getByTestId('byte-counter')).toHaveTextContent(`0/${defaultProps.maxBytes}`);
+      expect(mockClearDraft).toHaveBeenCalledTimes(1);
+      expect(mockClearDraft).toHaveBeenCalledWith(defaultProps.to);
     });
   });
 
-  it.skip('sends broadcast message if to is "broadcast" and updates state to Ack', async () => {
-    renderComponent({ to: 'broadcast', channel: 5, maxBytes: 256 });
-    const inputElement = screen.getByPlaceholderText('Enter Message') as HTMLInputElement;
-    fireEvent.change(inputElement, { target: { value: 'Broadcast message' } });
+  it('should trim whitespace before calling onSend', async () => {
+    renderComponent();
+    const inputElement = screen.getByPlaceholderText('Enter Message');
     const formElement = screen.getByRole('form');
+    const testMessageWithWhitespace = '  Trim me!  ';
+    const expectedTrimmedMessage = 'Trim me!';
+
+    fireEvent.change(inputElement, { target: { value: testMessageWithWhitespace } });
     fireEvent.submit(formElement);
 
     await waitFor(() => {
-      expect(mockSendText).toHaveBeenCalledWith('Broadcast message', 'broadcast', true, 5);
-      expect(mockSetMessageState).toHaveBeenCalledWith({
-        type: 'broadcast',
-        key: 123,
-        messageId: undefined,
-        newState: 'ack',
-      });
-      expect(mockClearDraft).toHaveBeenCalledWith('broadcast');
-      expect(inputElement.value).toBe('');
-      expect(screen.getByTestId('byte-counter')).toHaveTextContent('0/256');
+      expect(mockOnSend).toHaveBeenCalledTimes(1);
+      expect(mockOnSend).toHaveBeenCalledWith(expectedTrimmedMessage);
+      expect(mockClearDraft).toHaveBeenCalledWith(defaultProps.to);
     });
   });
 
-  it('updates state to Failed if sendText throws an error', async () => {
-    mockSendText.mockRejectedValue({ id: 456 });
-    renderComponent({ to: 3, channel: 1, maxBytes: 256 });
-    const inputElement = screen.getByPlaceholderText('Enter Message') as HTMLInputElement;
-    fireEvent.change(inputElement, { target: { value: 'Error message' } });
+  it('should not call onSend or clearDraft if input is empty on submit', async () => {
+    renderComponent();
+    const inputElement = screen.getByPlaceholderText('Enter Message');
     const formElement = screen.getByRole('form');
+
+    expect((inputElement as HTMLInputElement).value).toBe('');
+
     fireEvent.submit(formElement);
 
-    await waitFor(() => {
-      expect(mockSendText).toHaveBeenCalledWith('Error message', 3, true, 1);
-      expect(mockSetMessageState).toHaveBeenCalledWith({
-        type: 'direct',
-        key: 123,
-        messageId: 456,
-        newState: 'failed',
-      });
-      expect(mockClearDraft).toHaveBeenCalledWith(3);
-      expect(inputElement.value).toBe('');
-      expect(screen.getByTestId('byte-counter')).toHaveTextContent('0/256');
+    await act(async () => {
+      await new Promise(resolve => setTimeout(resolve, 50));
+    });
+
+    expect(mockOnSend).not.toHaveBeenCalled();
+    expect(mockClearDraft).not.toHaveBeenCalled();
+  });
+
+  it('should not call onSend or clearDraft if input contains only whitespace on submit', async () => {
+    renderComponent();
+    const inputElement = screen.getByTestId('message-input-field');
+    const formElement = screen.getByRole('form');
+    const whitespaceMessage = '   \t   ';
+
+    fireEvent.change(inputElement, { target: { value: whitespaceMessage } });
+    expect((inputElement as HTMLInputElement).value).toBe(whitespaceMessage);
+
+    fireEvent.submit(formElement);
+
+    await act(async () => {
+      await new Promise(resolve => setTimeout(resolve, 50));
     });
+
+    expect(mockOnSend).not.toHaveBeenCalled();
+    expect(mockClearDraft).not.toHaveBeenCalled();
+
+    expect((inputElement as HTMLInputElement).value).toBe(whitespaceMessage);
+  });
+
+  it('should work with broadcast destination for drafts', () => {
+    const broadcastDest: Types.Destination = 'broadcast';
+    mockGetDraft.mockImplementation((key) => key === broadcastDest ? 'Broadcast draft' : '');
+
+    renderComponent({ to: broadcastDest });
+
+    expect(mockGetDraft).toHaveBeenCalledWith(broadcastDest);
+    expect((screen.getByPlaceholderText('Enter Message') as HTMLInputElement).value).toBe('Broadcast draft');
+
+    const inputElement = screen.getByPlaceholderText('Enter Message');
+    const formElement = screen.getByRole('form');
+    const newMessage = 'New broadcast msg';
+
+    fireEvent.change(inputElement, { target: { value: newMessage } });
+    expect(mockSetDraft).toHaveBeenCalledWith(broadcastDest, newMessage);
+
+    fireEvent.submit(formElement);
+
+    expect(mockOnSend).toHaveBeenCalledWith(newMessage);
+    expect(mockClearDraft).toHaveBeenCalledWith(broadcastDest);
   });
 });

src/components/PageComponents/Messages/MessageInput.tsx

@@ -18,11 +18,12 @@ export const MessageInput = ({
 }: MessageInputProps) => {
   const { setDraft, getDraft, clearDraft } = useMessageStore();
 
-  const [localDraft, setLocalDraft] = useState(getDraft(to));
-  const [messageBytes, setMessageBytes] = useState(0);
-
   const calculateBytes = (text: string) => new Blob([text]).size;
 
+  const initialDraft = getDraft(to);
+  const [localDraft, setLocalDraft] = useState(initialDraft);
+  const [messageBytes, setMessageBytes] = useState(() => calculateBytes(initialDraft));
+
   const handleInputChange = (e: React.ChangeEvent<HTMLInputElement>) => {
     const newValue = e.target.value;
     const byteLength = calculateBytes(newValue);
@@ -37,6 +38,7 @@ export const MessageInput = ({
   const handleSubmit = (e: React.FormEvent) => {
     e.preventDefault();
     if (!localDraft.trim()) return;
+    // Reset bytes *before* sending (consider if onSend failure needs different handling)
     setMessageBytes(0);
 
     startTransition(() => {
@@ -67,8 +69,7 @@ export const MessageInput = ({
 
           <Button
             type="submit"
-            variant={"default"}
-          // className=" dark:text-slate-900 dark:hover:bg-slate-400 dark:hover:text-white"
+            variant="default"
           >
             <SendIcon size={16} />
           </Button>
@@ -76,4 +77,4 @@ export const MessageInput = ({
       </form>
     </div>
   );
-};
+};

src/components/PageComponents/Messages/MessageItem.tsx

@@ -88,18 +88,17 @@ export const MessageItem = ({ message }: MessageProps) => {
     return null;
   }, [getDevices, message.from]);
 
-  const { shortName, displayName } = useMemo(() => {
+  const { shortName, displayName, } = useMemo(() => {
     const fallbackName = message.from
     const longName = messageUser?.user?.longName;
     const shortName = messageUser?.user?.shortName ?? fallbackName;
     const displayName = longName || fallbackName;
-    return { shortName, displayName };
+    return { shortName, displayName, from: message.from };
   }, [messageUser, message.from]);
 
   const messageStatus = getMessageStatus(message.state);
   const messageText = message?.message ?? "";
   const messageDate = message?.date;
-  const isFailed = message.state === MessageState.Failed;
 
   const messageItemWrapperClass = cn(
     "group w-full py-2 relative list-none",

src/components/UI/Button.tsx

@@ -14,7 +14,7 @@ const buttonVariants = cva(
         success:
           "bg-green-500 text-white hover:bg-green-600 dark:hover:bg-green-600",
         outline:
-          "bg-transparent border border-slate-400 hover:text-slate-100 dark:hover:text-slate-300 dark:border-slate-400 dark:text-slate-100 ",
+          "bg-transparent border border-slate-400 hover:text-slate-400 dark:hover:text-slate-300 dark:border-slate-400 dark:text-slate-100 ",
         subtle:
           "bg-slate-100 text-slate-700 hover:bg-slate-200 dark:bg-slate-500 dark:text-white dark:hover:bg-slate-400",
         ghost:

src/components/UI/Dialog.tsx

@@ -61,6 +61,7 @@ const DialogClose = ({
 }: DialogPrimitive.DialogCloseProps & React.RefAttributes<HTMLButtonElement> & { className?: string }) => (
   <DialogPrimitive.Close
     aria-label="Close"
+    data-testid="dialog-close-button"
     className={cn(
       "absolute top-4 right-4 rounded-xs opacity-70 transition-opacity hover:opacity-100 focus:outline-hidden focus:ring-2 focus:ring-slate-400 focus:ring-offset-2 disabled:pointer-events-none data-[state=open]:bg-slate-100 dark:focus:ring-slate-400 dark:focus:ring-offset-slate-900",
       className,

src/core/stores/messageStore/index.ts

@@ -16,7 +16,7 @@ export enum MessageType {
   Broadcast = "broadcast",
 }
 
-function getConversationId(node1: NodeNum, node2: NodeNum): ConversationId {
+export function getConversationId(node1: NodeNum, node2: NodeNum): ConversationId {
   return [node1, node2].sort((a, b) => a - b).join(':');
 }
 
@@ -165,17 +165,17 @@ export const useMessageStore = create<MessageStore>()(
               const deleted = messageLog.delete(params.messageId);
 
               if (deleted) {
-                console.log(`Deleted message ${params.messageId} from ${params.type} chat ${parentKey}`);
+                console.log(`Deleted message ${params.messageId} from ${params.type} message ${parentKey}`);
                 // Clean up empty MessageLogMap and its entry in the parent map
                 if (messageLog.size === 0) {
                   parentMap.delete(parentKey);
-                  console.log(`Cleaned up empty chat entry for ${parentKey}`);
+                  console.log(`Cleaned up empty message entry for ${parentKey}`);
                 }
               } else {
                 console.warn(`Message ${params.messageId} not found in ${params.type} chat ${parentKey} for deletion.`);
               }
             } else {
-              console.warn(`Chat entry ${parentKey} not found for message deletion.`);
+              console.warn(`Message entry ${parentKey} not found for message deletion.`);
             }
           })
         );

src/core/stores/messageStore/messageStore.test.ts

@@ -3,14 +3,15 @@ import {
   useMessageStore,
   MessageType,
   MessageState,
-  type Message,
+  getConversationId,
 } from './index.ts';
+import type { ConversationId, ChannelId, MessageLogMap, Message } from './types.ts';
+import { Types } from '@meshtastic/core';
 
-let memoryStorage: Record<string, string> = {};
-
-vi.mock('./storage/indexDB.ts', () => {
+vi.mock('../storage/indexDB.ts', () => {
+  let memoryStorage: Record<string, string> = {};
   return {
-    zustandIndexDBStorage: {
+    storageWithMapSupport: {
       getItem: vi.fn(async (name: string): Promise<string | null> => {
         return memoryStorage[name] ?? null;
       }),
@@ -27,9 +28,7 @@ vi.mock('./storage/indexDB.ts', () => {
 const myNodeNum = 111;
 const otherNodeNum1 = 222;
 const otherNodeNum2 = 333;
-const broadcastChannel = 0;
-
-
+const broadcastChannel: ChannelId = 0;
 
 const directMessageToOther1: Message = {
   type: MessageType.Direct,
@@ -90,13 +89,23 @@ describe('useMessageStore', () => {
   const initialState = useMessageStore.getState();
 
   beforeEach(() => {
-    useMessageStore.setState(initialState, true);
+    useMessageStore.setState({
+      ...initialState,
+      messages: {
+        direct: new Map<ConversationId, MessageLogMap>(),
+        broadcast: new Map<ChannelId, MessageLogMap>(),
+      },
+      draft: new Map<Types.Destination, string>(),
+    }, true);
+
   });
 
   it('should have correct initial state', () => {
     const state = useMessageStore.getState();
-    expect(state.messages.direct).toEqual({});
-    expect(state.messages.broadcast).toEqual({});
+    expect(state.messages.direct).toBeInstanceOf(Map);
+    expect(state.messages.direct.size).toBe(0);
+    expect(state.messages.broadcast).toBeInstanceOf(Map);
+    expect(state.messages.broadcast.size).toBe(0);
     expect(state.draft).toBeInstanceOf(Map);
     expect(state.draft.size).toBe(0);
     expect(state.nodeNum).toBe(0);
@@ -117,23 +126,32 @@ describe('useMessageStore', () => {
   });
 
   describe('saveMessage', () => {
-    it('should save a direct message with correct structure', () => {
+    it('should save a direct message with correct Map structure', () => {
       useMessageStore.getState().saveMessage(directMessageToOther1);
       const state = useMessageStore.getState();
-      expect(state.messages.direct[myNodeNum]).toBeDefined();
-      expect(state.messages.direct[myNodeNum][otherNodeNum1]).toBeDefined();
-      expect(
-        state.messages.direct[myNodeNum][otherNodeNum1][directMessageToOther1.messageId],
-      ).toEqual(directMessageToOther1);
+      const conversationId = getConversationId(directMessageToOther1.from, directMessageToOther1.to);
+
+      // Check if the conversation Map exists
+      expect(state.messages.direct.has(conversationId)).toBe(true);
+      const conversationLog = state.messages.direct.get(conversationId);
+      // Check if the inner Map (MessageLogMap) exists and is a Map
+      expect(conversationLog).toBeInstanceOf(Map);
+      // Check if the message exists within the inner Map
+      expect(conversationLog?.has(directMessageToOther1.messageId)).toBe(true);
+      // Check the message content
+      expect(conversationLog?.get(directMessageToOther1.messageId)).toEqual(directMessageToOther1);
     });
 
-    it('should save a broadcast message with correct structure', () => {
+    it('should save a broadcast message with correct Map structure', () => {
       useMessageStore.getState().saveMessage(broadcastMessage1);
       const state = useMessageStore.getState();
-      expect(state.messages.broadcast[broadcastChannel]).toBeDefined();
-      expect(
-        state.messages.broadcast[broadcastChannel][broadcastMessage1.messageId],
-      ).toEqual(broadcastMessage1);
+      const channelId = broadcastMessage1.channel;
+
+      expect(state.messages.broadcast.has(channelId)).toBe(true);
+      const channelLog = state.messages.broadcast.get(channelId);
+      expect(channelLog).toBeInstanceOf(Map);
+      expect(channelLog?.has(broadcastMessage1.messageId)).toBe(true);
+      expect(channelLog?.get(broadcastMessage1.messageId)).toEqual(broadcastMessage1);
     });
 
     it('should save multiple messages correctly', () => {
@@ -143,12 +161,13 @@ describe('useMessageStore', () => {
 
       const state = useMessageStore.getState();
 
-      // Direct msg 1 (me -> other1)
-      expect(state.messages.direct[myNodeNum]?.[otherNodeNum1]?.[directMessageToOther1.messageId]).toEqual(directMessageToOther1);
-      // Direct msg 2 (other1 -> me)
-      expect(state.messages.direct[otherNodeNum1]?.[myNodeNum]?.[directMessageFromOther1.messageId]).toEqual(directMessageFromOther1);
-      // Broadcast msg 1
-      expect(state.messages.broadcast[broadcastChannel]?.[broadcastMessage1.messageId]).toEqual(broadcastMessage1);
+      const convId1 = getConversationId(myNodeNum, otherNodeNum1);
+      expect(state.messages.direct.get(convId1)?.get(directMessageToOther1.messageId)).toEqual(directMessageToOther1);
+
+      expect(state.messages.direct.get(convId1)?.get(directMessageFromOther1.messageId)).toEqual(directMessageFromOther1);
+
+      const channelId = broadcastMessage1.channel;
+      expect(state.messages.broadcast.get(channelId)?.get(broadcastMessage1.messageId)).toEqual(broadcastMessage1);
     });
   });
 
@@ -163,9 +182,9 @@ describe('useMessageStore', () => {
     });
 
     it('should return broadcast messages for a channel, sorted by date', () => {
-      const messages = useMessageStore.getState().getMessages(MessageType.Broadcast, {
-        myNodeNum: myNodeNum, // Not strictly needed for broadcast, but good practice
-        channel: broadcastChannel
+      const messages = useMessageStore.getState().getMessages({
+        type: MessageType.Broadcast,
+        channelId: broadcastChannel
       });
       expect(messages).toHaveLength(2);
       expect(messages[0]).toEqual(broadcastMessage1);
@@ -173,17 +192,18 @@ describe('useMessageStore', () => {
     });
 
     it('should return empty array for broadcast if channel has no messages', () => {
-      const messages = useMessageStore.getState().getMessages(MessageType.Broadcast, {
-        myNodeNum: myNodeNum,
-        channel: 99
+      const messages = useMessageStore.getState().getMessages({
+        type: MessageType.Broadcast,
+        channelId: Types.ChannelNumber.Channel1
       });
       expect(messages).toEqual([]);
     });
 
-    it('should return combined direct messages for a specific chat (pair), sorted by date', () => {
-      const messages = useMessageStore.getState().getMessages(MessageType.Direct, {
-        myNodeNum: myNodeNum,
-        otherNodeNum: otherNodeNum1
+    it('should return combined direct messages for a specific chat pair, sorted by date', () => {
+      const messages = useMessageStore.getState().getMessages({
+        type: MessageType.Direct,
+        nodeA: myNodeNum,
+        nodeB: otherNodeNum1
       });
       expect(messages).toHaveLength(2);
       expect(messages[0]).toEqual(directMessageToOther1);
@@ -191,31 +211,23 @@ describe('useMessageStore', () => {
     });
 
     it('should return only relevant direct messages for a different chat pair', () => {
-      const messages = useMessageStore.getState().getMessages(MessageType.Direct, {
-        myNodeNum: myNodeNum,
-        otherNodeNum: otherNodeNum2
+      const messages = useMessageStore.getState().getMessages({
+        type: MessageType.Direct,
+        nodeA: myNodeNum,
+        nodeB: otherNodeNum2
       });
       expect(messages).toHaveLength(1);
       expect(messages[0]).toEqual(directMessageToOther2);
     });
 
     it('should return empty array for direct chat if no messages exist', () => {
-      const messages = useMessageStore.getState().getMessages(MessageType.Direct, {
-        myNodeNum: myNodeNum,
-        otherNodeNum: 999
+      const messages = useMessageStore.getState().getMessages({
+        type: MessageType.Direct,
+        nodeA: myNodeNum,
+        nodeB: 999
       });
       expect(messages).toEqual([]);
     });
-
-    it('should return combined direct messages when myNodeNum and otherNodeNum are provided', () => {
-      const messages = useMessageStore.getState().getMessages(MessageType.Direct, {
-        myNodeNum: myNodeNum,         // Keep this
-        otherNodeNum: otherNodeNum1
-      });
-      expect(messages).toHaveLength(2);
-      expect(messages[0]).toEqual(directMessageToOther1);
-      expect(messages[1]).toEqual(directMessageFromOther1);
-    });
   });
 
   describe('setMessageState', () => {
@@ -226,119 +238,219 @@ describe('useMessageStore', () => {
       useMessageStore.getState().saveMessage(broadcastMessage1);
     });
 
-    it('should update state for a direct message sent BY ME', () => {
+    it('should update state for a direct message', () => {
       useMessageStore.getState().setMessageState({
         type: MessageType.Direct,
-        key: otherNodeNum1,
+        nodeA: directMessageToOther1.from,
+        nodeB: directMessageToOther1.to,
         messageId: directMessageToOther1.messageId,
         newState: MessageState.Ack,
       });
-      const message = useMessageStore.getState().messages.direct[myNodeNum]?.[otherNodeNum1]?.[directMessageToOther1.messageId];
+      const conversationId = getConversationId(directMessageToOther1.from, directMessageToOther1.to);
+      const message = useMessageStore.getState().messages.direct.get(conversationId)?.get(directMessageToOther1.messageId);
       expect(message?.state).toBe(MessageState.Ack);
     });
 
-    it('should update state for a direct message received FROM OTHER', () => {
+    it('should update state for another direct message in the same conversation', () => {
       useMessageStore.getState().setMessageState({
         type: MessageType.Direct,
-        key: otherNodeNum1,
+        nodeA: directMessageFromOther1.from,
+        nodeB: directMessageFromOther1.to,
         messageId: directMessageFromOther1.messageId,
         newState: MessageState.Failed,
       });
-      const message = useMessageStore.getState().messages.direct[otherNodeNum1]?.[myNodeNum]?.[directMessageFromOther1.messageId];
+      const conversationId = getConversationId(directMessageFromOther1.from, directMessageFromOther1.to);
+      const message = useMessageStore.getState().messages.direct.get(conversationId)?.get(directMessageFromOther1.messageId);
       expect(message?.state).toBe(MessageState.Failed);
     });
 
     it('should update state for a broadcast message', () => {
       useMessageStore.getState().setMessageState({
         type: MessageType.Broadcast,
-        key: broadcastChannel,
+        channelId: broadcastChannel,
         messageId: broadcastMessage1.messageId,
         newState: MessageState.Ack,
       });
-      const message = useMessageStore.getState().messages.broadcast[broadcastChannel]?.[broadcastMessage1.messageId];
+      const message = useMessageStore.getState().messages.broadcast.get(broadcastChannel)?.get(broadcastMessage1.messageId);
       expect(message?.state).toBe(MessageState.Ack);
     });
 
-    it('should warn if message is not found', () => {
+    it('should warn if message is not found (direct)', () => {
       const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => { });
       useMessageStore.getState().setMessageState({
         type: MessageType.Direct,
-        key: otherNodeNum1,
+        nodeA: myNodeNum,
+        nodeB: otherNodeNum1,
         messageId: 999,
         newState: MessageState.Ack,
       });
-      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('Message not found for state update'));
+      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('Message or conversation/channel not found for state update'));
       warnSpy.mockRestore();
     });
-  });
 
+    it('should warn if message is not found (broadcast)', () => {
+      const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => { });
+      useMessageStore.getState().setMessageState({
+        type: MessageType.Broadcast,
+        channelId: broadcastChannel,
+        messageId: 999,
+        newState: MessageState.Ack,
+      });
+      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('Message or conversation/channel not found for state update'));
+      warnSpy.mockRestore();
+    });
+
+    it('should warn if conversation/channel is not found', () => {
+      const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => { });
+      useMessageStore.getState().setMessageState({
+        type: MessageType.Direct,
+        nodeA: myNodeNum,
+        nodeB: 998,
+        messageId: 101,
+        newState: MessageState.Ack,
+      });
+      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('Message or conversation/channel not found for state update'));
+      warnSpy.mockRestore();
+    });
+  });
 
   describe('clearMessageByMessageId', () => {
+    const extraDirectMessageId = 1011;
     beforeEach(() => {
       useMessageStore.getState().setNodeNum(myNodeNum);
       useMessageStore.getState().saveMessage(directMessageToOther1);
       useMessageStore.getState().saveMessage(directMessageFromOther1);
       useMessageStore.getState().saveMessage(broadcastMessage1);
-      useMessageStore.getState().saveMessage({ ...directMessageToOther1, messageId: 1011, date: Date.now() + 50 });
+      useMessageStore.getState().saveMessage({ ...directMessageToOther1, messageId: extraDirectMessageId, date: Date.now() + 50 });
     });
 
-    it('should delete a specific direct message (sent by me)', () => {
+    it('should delete a specific direct message', () => {
       const messageIdToDelete = directMessageToOther1.messageId;
+      const nodeA = directMessageToOther1.from;
+      const nodeB = directMessageToOther1.to;
+      const conversationId = getConversationId(nodeA, nodeB);
+
       useMessageStore.getState().clearMessageByMessageId({
         type: MessageType.Direct,
-        from: myNodeNum,
-        to: otherNodeNum1,
+        nodeA: nodeA,
+        nodeB: nodeB,
         messageId: messageIdToDelete
       });
+
       const state = useMessageStore.getState();
-      expect(state.messages.direct[myNodeNum]?.[otherNodeNum1]?.[messageIdToDelete]).toBeUndefined();
-      expect(state.messages.direct[myNodeNum]?.[otherNodeNum1]?.[1011]).toBeDefined();
-      expect(state.messages.direct[otherNodeNum1]?.[myNodeNum]?.[directMessageFromOther1.messageId]).toBeDefined();
+      const conversationLog = state.messages.direct.get(conversationId);
+      expect(conversationLog?.has(messageIdToDelete)).toBe(false);
+      expect(conversationLog?.has(extraDirectMessageId)).toBe(true);
+      expect(conversationLog?.has(directMessageFromOther1.messageId)).toBe(true);
+      expect(state.messages.direct.has(conversationId)).toBe(true);
+
     });
 
-    it('should delete a specific direct message (sent by other)', () => {
+    it('should delete another specific direct message', () => {
       const messageIdToDelete = directMessageFromOther1.messageId;
+      const nodeA = directMessageFromOther1.from;
+      const nodeB = directMessageFromOther1.to;
+      const conversationId = getConversationId(nodeA, nodeB);
+
       useMessageStore.getState().clearMessageByMessageId({
         type: MessageType.Direct,
-        from: otherNodeNum1,
-        to: myNodeNum,
+        nodeA: nodeA,
+        nodeB: nodeB,
         messageId: messageIdToDelete
       });
+
       const state = useMessageStore.getState();
-      expect(state.messages.direct[otherNodeNum1]?.[myNodeNum]?.[messageIdToDelete]).toBeUndefined();
-      expect(state.messages.direct[myNodeNum]?.[otherNodeNum1]?.[directMessageToOther1.messageId]).toBeDefined();
-      expect(state.messages.direct[myNodeNum]?.[otherNodeNum1]?.[1011]).toBeDefined();
+      const conversationLog = state.messages.direct.get(conversationId);
+      expect(conversationLog?.has(messageIdToDelete)).toBe(false);
+      expect(conversationLog?.has(directMessageToOther1.messageId)).toBe(true);
+      expect(conversationLog?.has(extraDirectMessageId)).toBe(true);
     });
 
+
     it('should delete a specific broadcast message', () => {
       const messageIdToDelete = broadcastMessage1.messageId;
+      const channelId = broadcastMessage1.channel;
+
       useMessageStore.getState().clearMessageByMessageId({
         type: MessageType.Broadcast,
-        channel: broadcastChannel,
+        channelId: channelId,
         messageId: messageIdToDelete
       });
+
+      const state = useMessageStore.getState();
+      expect(state.messages.broadcast.get(channelId)?.get(messageIdToDelete)).toBeUndefined();
+    });
+
+    it('should clean up empty conversation/channel Maps', () => {
+      const directConvId = getConversationId(directMessageFromOther1.from, directMessageFromOther1.to);
+      const broadcastChanId = broadcastMessage1.channel;
+
+      useMessageStore.getState().clearMessageByMessageId({ type: MessageType.Direct, nodeA: directMessageToOther1.from, nodeB: directMessageToOther1.to, messageId: directMessageToOther1.messageId });
+      useMessageStore.getState().clearMessageByMessageId({ type: MessageType.Direct, nodeA: directMessageFromOther1.from, nodeB: directMessageFromOther1.to, messageId: directMessageFromOther1.messageId });
+      useMessageStore.getState().clearMessageByMessageId({ type: MessageType.Direct, nodeA: directMessageToOther1.from, nodeB: directMessageToOther1.to, messageId: extraDirectMessageId });
+
+      expect(useMessageStore.getState().messages.direct.has(directConvId)).toBe(false);
+
+      useMessageStore.getState().clearMessageByMessageId({ type: MessageType.Broadcast, channelId: broadcastChanId, messageId: broadcastMessage1.messageId });
+
+      expect(useMessageStore.getState().messages.broadcast.has(broadcastChanId)).toBe(false);
+    });
+
+    it('should not error when trying to delete non-existent message', () => {
+      const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => { });
+      const conversationId = getConversationId(myNodeNum, otherNodeNum1);
+
+      expect(() => {
+        useMessageStore.getState().clearMessageByMessageId({
+          type: MessageType.Direct,
+          nodeA: myNodeNum,
+          nodeB: otherNodeNum1,
+          messageId: 9999
+        });
+      }).not.toThrow();
+
       const state = useMessageStore.getState();
-      expect(state.messages.broadcast[broadcastChannel]?.[messageIdToDelete]).toBeUndefined();
+      const conversationLog = state.messages.direct.get(conversationId);
+      expect(conversationLog?.size).toBe(3); // 101, 102, 1011
+      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('not found in direct chat'));
+
+      warnSpy.mockRestore();
     });
 
-    it('should clean up empty to/from/channel objects', () => {
-      useMessageStore.getState().clearMessageByMessageId({ type: MessageType.Direct, from: otherNodeNum1, to: myNodeNum, messageId: directMessageFromOther1.messageId });
-      expect(useMessageStore.getState().messages.direct[otherNodeNum1]?.[myNodeNum]).toBeUndefined(); // Recipient level removed
-      expect(useMessageStore.getState().messages.direct[otherNodeNum1]).toBeUndefined(); // Sender level removed
+    it('should not error when trying to delete from non-existent conversation/channel', () => {
+      const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => { });
+      expect(() => {
+        useMessageStore.getState().clearMessageByMessageId({
+          type: MessageType.Direct,
+          nodeA: myNodeNum,
+          nodeB: 9998,
+          messageId: 101
+        });
+      }).not.toThrow();
+
+      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("Message entry"));
 
-      useMessageStore.getState().clearMessageByMessageId({ type: MessageType.Broadcast, channel: broadcastChannel, messageId: broadcastMessage1.messageId });
-      expect(useMessageStore.getState().messages.broadcast[broadcastChannel]).toBeUndefined(); // Channel level removed
+      expect(warnSpy).toHaveBeenCalledTimes(1);
+
+      warnSpy.mockRestore();
     });
   });
 
   describe('Drafts', () => {
-    const draftKey = otherNodeNum1;
+    const draftKeyDirect = otherNodeNum1;
+    const draftKeyBroadcast = broadcastChannel;
     const draftMessage = 'This is a draft';
 
-    it('should set and get a draft', () => {
-      useMessageStore.getState().setDraft(draftKey, draftMessage);
-      expect(useMessageStore.getState().draft.get(draftKey)).toBe(draftMessage);
-      expect(useMessageStore.getState().getDraft(draftKey)).toBe(draftMessage);
+    it('should set and get a draft for direct chat', () => {
+      useMessageStore.getState().setDraft(draftKeyDirect, draftMessage);
+      expect(useMessageStore.getState().draft.get(draftKeyDirect)).toBe(draftMessage);
+      expect(useMessageStore.getState().getDraft(draftKeyDirect)).toBe(draftMessage);
+    });
+
+    it('should set and get a draft for broadcast chat', () => {
+      useMessageStore.getState().setDraft(draftKeyBroadcast, draftMessage);
+      expect(useMessageStore.getState().draft.get(draftKeyBroadcast)).toBe(draftMessage);
+      expect(useMessageStore.getState().getDraft(draftKeyBroadcast)).toBe(draftMessage);
     });
 
     it('should return empty string for non-existent draft', () => {
@@ -346,26 +458,29 @@ describe('useMessageStore', () => {
     });
 
     it('should clear a draft', () => {
-      useMessageStore.getState().setDraft(draftKey, draftMessage);
-      expect(useMessageStore.getState().draft.has(draftKey)).toBe(true);
-      useMessageStore.getState().clearDraft(draftKey);
-      expect(useMessageStore.getState().draft.has(draftKey)).toBe(false);
-      expect(useMessageStore.getState().getDraft(draftKey)).toBe('');
+      useMessageStore.getState().setDraft(draftKeyDirect, draftMessage);
+      expect(useMessageStore.getState().draft.has(draftKeyDirect)).toBe(true);
+      useMessageStore.getState().clearDraft(draftKeyDirect);
+      expect(useMessageStore.getState().draft.has(draftKeyDirect)).toBe(false);
+      expect(useMessageStore.getState().getDraft(draftKeyDirect)).toBe('');
     });
   });
 
   describe('deleteAllMessages', () => {
-    it('should clear all direct and broadcast messages', () => {
+    it('should clear all direct and broadcast messages, leaving empty Maps', () => {
       useMessageStore.getState().saveMessage(directMessageToOther1);
       useMessageStore.getState().saveMessage(broadcastMessage1);
-      expect(Object.keys(useMessageStore.getState().messages.direct).length).toBeGreaterThan(0);
-      expect(Object.keys(useMessageStore.getState().messages.broadcast).length).toBeGreaterThan(0);
+
+      expect(useMessageStore.getState().messages.direct.size).toBeGreaterThan(0);
+      expect(useMessageStore.getState().messages.broadcast.size).toBeGreaterThan(0);
 
       useMessageStore.getState().deleteAllMessages();
 
-      expect(useMessageStore.getState().messages.direct).toEqual({});
-      expect(useMessageStore.getState().messages.broadcast).toEqual({});
+      const state = useMessageStore.getState();
+      expect(state.messages.direct).toBeInstanceOf(Map);
+      expect(state.messages.direct.size).toBe(0);
+      expect(state.messages.broadcast).toBeInstanceOf(Map);
+      expect(state.messages.broadcast.size).toBe(0);
     });
   });
-
-});
+});