diff --git a/ui/src/components/messages/MessageState.tsx b/ui/src/components/messages/MessageState.tsx
index 638b9dc..d0caee0 100644
--- a/ui/src/components/messages/MessageState.tsx
+++ b/ui/src/components/messages/MessageState.tsx
@@ -1,32 +1,36 @@
 import React, { createContext, useContext, useState, ReactNode } from 'react'
 import { Message } from './SubscriptionsHook'
 
-export type MessageStateType = {
+type MessageContextType = {
   messages: Message[]
   setMessages: React.Dispatch<React.SetStateAction<Message[]>>
-  add(message: Message): void
-  removeKey(key: string): void
+  add: (message: Message) => void
+  removeKey: (key: string) => void
 }
 
-const MessageContext = createContext<MessageStateType | undefined>(undefined)
+const MessageContext = createContext<MessageContextType | undefined>(undefined)
 
-interface MessageProviderProps {
+export const useMessageState = () => {
+  const context = useContext(MessageContext)
+  if (context === undefined) {
+    throw new Error('useMessageState must be used within a MessageProvider')
+  }
+  return context
+}
+
+type MessageProviderProps = {
   children: ReactNode
 }
 
-export const MessageProvider: React.FC<MessageProviderProps> = ({ children }) => {
+export const MessageProvider = ({ children }: MessageProviderProps) => {
   const [messages, setMessages] = useState<Message[]>([])
 
   const add = (message: Message) => {
-    setMessages(prevMessages => {
-      const newMessages = prevMessages.filter(msg => msg.key != message.key)
-      newMessages.push(message)
-      return newMessages
-    })
+    setMessages((prevMessages) => [...prevMessages, message])
   }
 
   const removeKey = (key: string) => {
-    setMessages(prevMessages => prevMessages.filter(msg => msg.key != key))
+    setMessages((prevMessages) => prevMessages.filter((msg) => msg.key !== key))
   }
 
   return (
@@ -35,11 +39,3 @@ export const MessageProvider: React.FC<MessageProviderProps> = ({ children }) =>
     </MessageContext.Provider>
   )
 }
-
-export const useMessageState = (): MessageStateType => {
-  const context = useContext(MessageContext)
-  if (!context) {
-    throw new Error('useMessageState must be used within a MessageProvider')
-  }
-  return context
-}
diff --git a/ui/src/components/messages/Messages.test.tsx b/ui/src/components/messages/Messages.test.tsx
index c3c28b2..d68054e 100644
--- a/ui/src/components/messages/Messages.test.tsx
+++ b/ui/src/components/messages/Messages.test.tsx
@@ -21,7 +21,7 @@ const MockSubscriptionsHook = ({ messages, setMessages }: any) => {
 
 vi.mock('./SubscriptionsHook', () => ({
   __esModule: true,
-  SubscriptionsHook: MockSubscriptionsHook,
+  default: MockSubscriptionsHook,
 }))
 
 const messages = [
@@ -76,16 +76,16 @@ describe('Messages Component', () => {
       </MessageProvider>
     )
 
-    const neutralMessage = screen.getByText('This is a neutral message.')
-    const positiveMessage = screen.getByText('This is a positive message.')
-    const negativeMessage = screen.getByText('This is a negative message.')
+    const neutralMessage = screen.getByText('This is a neutral message.').closest('div')
+    const positiveMessage = screen.getByText('This is a positive message.').closest('div')
+    const negativeMessage = screen.getByText('This is a negative message.').closest('div')
 
-    expect(neutralMessage.closest('div')).toHaveClass('bg-white')
-    expect(neutralMessage.closest('div')).toHaveClass('dark:bg-dark-bg2')
-    expect(positiveMessage.closest('div')).toHaveClass('bg-green-100')
-    expect(positiveMessage.closest('div')).toHaveClass('dark:bg-green-900')
-    expect(negativeMessage.closest('div')).toHaveClass('bg-red-100')
-    expect(negativeMessage.closest('div')).toHaveClass('dark:bg-red-900')
+    expect(neutralMessage?.parentElement).toHaveClass('bg-white')
+    expect(neutralMessage?.parentElement).toHaveClass('dark:bg-dark-bg2')
+    expect(positiveMessage?.parentElement).toHaveClass('bg-green-100')
+    expect(positiveMessage?.parentElement).toHaveClass('dark:bg-green-900')
+    expect(negativeMessage?.parentElement).toHaveClass('bg-red-100')
+    expect(negativeMessage?.parentElement).toHaveClass('dark:bg-red-900')
   })
 
   test('renders messages with overflow and scroll behavior', () => {
@@ -114,7 +114,7 @@ describe('Messages Component', () => {
       </MessageProvider>
     )
 
-    const progressBar = container.querySelector('div[style*="width: 50%"]')
+    const progressBar = container.querySelector('div[role="progressbar"][style*="width: 50%"]')
     expect(progressBar).toBeInTheDocument()
     expect(progressBar).toHaveStyle('background-color: #fbbf24')
   })
