2 files changed, 48 insertions, 13 deletions
diff --git a/src/services/api/api.service.js b/src/services/api/api.service.js
index 7f27564f..c6818df4 100644
--- a/src/services/api/api.service.js
+++ b/src/services/api/api.service.js
@@ -953,8 +953,52 @@ const MASTODON_STREAMING_EVENTS = new Set([
   'filters_changed'
 ])
 
+// A thin wrapper around WebSocket API that allows adding a pre-processor to it
+// Uses EventTarget and a CustomEvent to proxy events
+export const ProcessedWS = ({
+  url,
+  preprocessor = handleMastoWS,
+  id = 'Unknown'
+}) => {
+  const eventTarget = new EventTarget()
+  const socket = new WebSocket(url)
+  if (!socket) throw new Error(`Failed to create socket ${id}`)
+  const proxy = (original, eventName, processor = a => a) => {
+    original.addEventListener(eventName, (eventData) => {
+      eventTarget.dispatchEvent(new CustomEvent(
+        eventName,
+        { detail: processor(eventData) }
+      ))
+    })
+  }
+  socket.addEventListener('open', (wsEvent) => {
+    console.debug(`[WS][${id}] Socket connected`, wsEvent)
+  })
+  socket.addEventListener('error', (wsEvent) => {
+    console.debug(`[WS][${id}] Socket errored`, wsEvent)
+  })
+  socket.addEventListener('close', (wsEvent) => {
+    console.debug(
+      `[WS][${id}] Socket disconnected with code ${wsEvent.code}`,
+      wsEvent
+    )
+  })
+  socket.addEventListener('message', (wsEvent) => {
+    console.debug(
+      `[WS][${id}] Message received`,
+      wsEvent
+    )
+  })
+
+  proxy(socket, 'open')
+  proxy(socket, 'close')
+  proxy(socket, 'message', preprocessor)
+  proxy(socket, 'error')
+
+  return eventTarget
+}
+
 export const handleMastoWS = (wsEvent) => {
-  console.debug('Event', wsEvent)
   const { data } = wsEvent
   if (!data) return
   const parsedEvent = JSON.parse(data)
diff --git a/src/services/backend_interactor_service/backend_interactor_service.js b/src/services/backend_interactor_service/backend_interactor_service.js
index 33b79a40..b7372ed0 100644
--- a/src/services/backend_interactor_service/backend_interactor_service.js
+++ b/src/services/backend_interactor_service/backend_interactor_service.js
@@ -1,4 +1,4 @@
-import apiService, { getMastodonSocketURI, handleMastoWS } from '../api/api.service.js'
+import apiService, { getMastodonSocketURI, ProcessedWS } from '../api/api.service.js'
 import timelineFetcherService from '../timeline_fetcher/timeline_fetcher.service.js'
 import notificationsFetcher from '../notifications_fetcher/notifications_fetcher.service.js'
 import followRequestFetcher from '../../services/follow_request_fetcher/follow_request_fetcher.service'
@@ -20,19 +20,10 @@ const backendInteractorService = credentials => ({
     return followRequestFetcher.startFetching({ store, credentials })
   },
 
-  startUserSocket ({ store, onMessage }) {
+  startUserSocket ({ store }) {
     const serv = store.rootState.instance.server.replace('http', 'ws')
     const url = serv + getMastodonSocketURI({ credentials, stream: 'user' })
-    const socket = new WebSocket(url)
-    console.debug('Socket created:', socket)
-    if (socket) {
-      socket.addEventListener('open', (wsEvent) => console.debug('MastoAPI User WebSocket connection established'))
-      socket.addEventListener('message', (wsEvent) => onMessage(handleMastoWS(wsEvent)))
-      socket.addEventListener('error', (error) => console.error('MastoApi User WebSocket Error:', error))
-      return socket
-    } else {
-      throw new Error('failed to connect to socket')
-    }
+    return ProcessedWS({ url, id: 'User' })
   },
 
   ...Object.entries(apiService).reduce((acc, [key, func]) => {