feat(useWebSocket): pass the retried to the autoReconnect.retries (#4604)

9romise · OrbisK · web-flow · commit 73e6db27cee3 · 2025-03-05T03:59:52.000Z
Co-authored-by: Robin &lt;robin.kehl@singular-it.de&gt;
diff --git a/packages/core/useWebSocket/index.browser.test.ts b/packages/core/useWebSocket/index.browser.test.ts
@@ -0,0 +1,92 @@
+import { useWebSocket } from '@vueuse/core'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+describe('useWebSocket', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+  })
+  afterEach(() => {
+    vi.restoreAllMocks()
+  })
+
+  describe('autoReconnect', () => {
+    const AUTO_RECONNECT_DELAY = 1000
+
+    it('should reconnect if autoReconnect is true', () => {
+      const { ws, status } = useWebSocket('ws://localhost', {
+        autoReconnect: true,
+      })
+
+      ws.value?.onopen?.(new Event('open'))
+
+      ws.value?.onclose?.(new CloseEvent('close'))
+      expect(status.value).toBe('CLOSED')
+      vi.advanceTimersByTime(AUTO_RECONNECT_DELAY)
+
+      expect(status.value).toBe('CONNECTING')
+    })
+
+    it('should not reconnect if autoReconnect is false', () => {
+      const { ws, status } = useWebSocket('ws://localhost', {
+        autoReconnect: false,
+      })
+
+      ws.value?.onopen?.(new Event('open'))
+
+      ws.value?.onclose?.(new CloseEvent('close'))
+      expect(status.value).toBe('CLOSED')
+      vi.advanceTimersByTime(AUTO_RECONNECT_DELAY)
+
+      expect(status.value).toBe('CLOSED')
+    })
+
+    it('should call onFailed if autoReconnect is not false', () => {
+      const onFailed = vi.fn()
+      const { ws, status } = useWebSocket('ws://localhost', {
+        autoReconnect: {
+          retries: () => false,
+          onFailed,
+        },
+      })
+
+      ws.value?.onopen?.(new Event('open'))
+
+      ws.value?.onclose?.(new CloseEvent('close'))
+
+      expect(onFailed).toHaveBeenCalled()
+      expect(status.value).toBe('CLOSED')
+    })
+
+    it('should reconnect if autoReconnect.retries is number', () => {
+      const { ws, status } = useWebSocket('ws://localhost', {
+        autoReconnect: {
+          retries: 2,
+        },
+      })
+
+      ws.value?.onopen?.(new Event('open'))
+
+      ws.value?.onclose?.(new CloseEvent('close'))
+      expect(status.value).toBe('CLOSED')
+      vi.advanceTimersByTime(AUTO_RECONNECT_DELAY)
+
+      expect(status.value).toBe('CONNECTING')
+    })
+
+    it('should reconnect if autoReconnect.retries is function', () => {
+      const { ws, status } = useWebSocket('ws://localhost', {
+        autoReconnect: {
+          retries: retried => retried < 1,
+        },
+      })
+
+      ws.value?.onopen?.(new Event('open'))
+
+      ws.value?.onclose?.(new CloseEvent('close'))
+      expect(status.value).toBe('CLOSED')
+      vi.advanceTimersByTime(AUTO_RECONNECT_DELAY)
+
+      expect(status.value).toBe('CONNECTING')
+    })
+  })
+})
diff --git a/packages/core/useWebSocket/index.ts b/packages/core/useWebSocket/index.ts
@@ -61,7 +61,7 @@ export interface UseWebSocketOptions {
      *
      * @default -1
      */
-    retries?: number | (() => boolean)
+    retries?: number | ((retried: number) => boolean)
 
     /**
      * Delay for reconnect, in milliseconds
@@ -256,13 +256,14 @@ export function useWebSocket<Data = any>(
           onFailed,
         } = resolveNestedOptions(options.autoReconnect)
 
-        if (typeof retries === 'number' && (retries < 0 || retried < retries)) {
+        const checkRetires = typeof retries === 'function'
+          ? retries
+          : () => typeof retries === 'number' && (retries < 0 || retried < retries)
+
+        if (checkRetires(retried)) {
           retried += 1
           retryTimeout = setTimeout(_init, delay)
         }
-        else if (typeof retries === 'function' && retries()) {
-          retryTimeout = setTimeout(_init, delay)
-        }
         else {
           onFailed?.()
         }
