fix: bug where heartbeat wasn't working (#18)

* fix: bug where heartbeat wasn't working

* chore: improve types

* style: format code

* test: add new test for heartbeat, update another tests

* chore: remove unused code from vite.config

src/close.ts

@@ -1,10 +1,15 @@
 import { heartbeatStop } from './heartbeat'
+import type { WSGOConfig } from './types'
 
-export function close(ws: WebSocket, ...[code = 1000, reason]: Parameters<WebSocket['close']>): void {
-  if (ws === undefined) return
-
-  // stop heartbeat interval
-  heartbeatStop()
+export function close(
+  ws: WebSocket,
+  _config: WSGOConfig,
+  ...[code = 1000, reason]: Parameters<WebSocket['close']>
+): void {
+  if (_config.heartbeat) {
+    // stop heartbeat interval
+    heartbeatStop()
+  }
 
   // close websocket connection
   ws.close(code, reason)

src/heartbeat/index.ts

@@ -1,28 +1,32 @@
 import { send } from '../send'
-import type { WSGOConfig } from '../types'
+import type { WSGOConfig, WSGOMessage } from '../types'
 
 const heartbeatMessage = 'ping'
 const heartbeatInterval = 1000
+const heartbeatTimeout = 1000
 
-let heartbeatTimeout: ReturnType<typeof setTimeout> | undefined
+let heartbeatTimeoutWait: ReturnType<typeof setTimeout> | undefined
 
 export function heartbeatStart(ws: WebSocket, _config: WSGOConfig): void {
-  heartbeatStop()
-
-  heartbeatTimeout = setTimeout(() => {
+  setTimeout(() => {
     send(ws, _config, heartbeatMessage, heartbeatMessage)
+    heartbeatTimeoutWait = setTimeout(() => {
+      // open()
+      ws.close()
+    }, heartbeatInterval + heartbeatTimeout)
   }, heartbeatInterval)
 }
 
-export function heartbeatStop(): void {
-  clearTimeout(heartbeatTimeout)
-  heartbeatTimeout = undefined
-}
-
-export function listenHeartbeat(ws: WebSocket, _config: WSGOConfig, e: MessageEvent<any>): void {
-  if (e.data === heartbeatMessage) {
+export function heartbeatListen(ws: WebSocket, _config: WSGOConfig, e: WSGOMessage<any>): void {
+  if (e.event === 'pong') {
+    heartbeatStop()
     heartbeatStart(ws, _config)
     // eslint-disable-next-line no-useless-return
     return
   }
 }
+
+export function heartbeatStop(): void {
+  clearTimeout(heartbeatTimeoutWait)
+  heartbeatTimeoutWait = undefined
+}

src/index.ts

@@ -1,12 +1,12 @@
-import { type WSGOEventName, type WSGOConfig, type WSGOSubscriptions } from './types'
+import type { WSGOEventName, WSGOConfig, WSGOSubscriptions, WSGOMessage } from './types'
 
 import { open } from './open'
 import { close } from './close'
 import { send } from './send'
 import type { WSGOSendData } from './send/types'
 import { subscribe } from './subscribe'
 import type { WSGOSubscribeCallback } from './subscribe/types'
-import { heartbeatStart, heartbeatStop, listenHeartbeat } from './heartbeat'
+import { heartbeatStart, heartbeatStop, heartbeatListen } from './heartbeat'
 
 /** Method allows you create new WebSocket connection */
 export default function create(
@@ -61,7 +61,7 @@ export default function create(
     close: () => {
       if (ws === undefined) return
 
-      close(ws)
+      close(ws, _config)
     },
     send: (...args) => {
       if (ws === undefined) return
@@ -81,20 +81,8 @@ function _listen(ws: WebSocket, subscriptions: WSGOSubscriptions, _config: WSGOC
     heartbeatStart(ws, _config)
   }
 
-  ws.onclose = (ev) => {
-    _config.onDisconnected?.(ws, ev)
-
-    heartbeatStop()
-  }
-
-  ws.onerror = (ev) => {
-    _config.onError?.(ws, ev)
-  }
-
   ws.onmessage = (e: MessageEvent<any>): any => {
-    listenHeartbeat(ws, _config, e)
-
-    let message
+    let message: WSGOMessage<any>
 
     try {
       message = JSON.parse(e.data)
@@ -106,8 +94,20 @@ function _listen(ws: WebSocket, subscriptions: WSGOSubscriptions, _config: WSGOC
       return
     }
 
+    heartbeatListen(ws, _config, message)
+
     if (message.event in subscriptions) {
       subscriptions[message.event](message)
     }
   }
+
+  ws.onerror = (ev) => {
+    _config.onError?.(ws, ev)
+  }
+
+  ws.onclose = (ev) => {
+    _config.onDisconnected?.(ws, ev)
+
+    heartbeatStop()
+  }
 }

src/subscribe/types.d.ts

@@ -0,0 +1,8 @@
+import type { WSGOMessage } from '../types'
+
+export interface WSGOSubscribeResponse<T = any> extends WSGOMessage<T> {
+  /** Time when the client received the event */
+  timeReceived: number
+}
+
+export type WSGOSubscribeCallback<T> = (message: WSGOSubscribeResponse<T>) => any

src/types/index.ts → src/types.d.ts

@@ -1,4 +1,4 @@
-import type { WSGOHeartbeat } from '../heartbeat/types'
+import type { WSGOHeartbeat } from './heartbeat/types'
 
 export type WSGOEventName = string
 
@@ -13,3 +13,12 @@ export interface WSGOConfig {
   immediate: boolean
   heartbeat: WSGOHeartbeat
 }
+
+export interface WSGOMessage<T> {
+  /** Event name */
+  event: WSGOEventName
+  /** Event data */
+  data: T
+  /** Time when the server sent the event */
+  timeSended: number
+}

test/close.test.ts

@@ -1,26 +1,21 @@
-import { describe, it, expect, beforeAll, afterAll } from 'vitest'
+import { describe, it, expect, beforeEach, afterEach } from 'vitest'
 import WSGO from '../src/index'
-import ws from 'ws'
 import { createMockWSServer } from './utils'
 
 describe('close', () => {
-  let port: number = 0
-  let server: ws.Server
+  let mockWSServer: ReturnType<typeof createMockWSServer>
 
-  beforeAll(() => {
-    const mockWSServer = createMockWSServer(port)
-
-    server = mockWSServer.server
-    port = mockWSServer.port
+  beforeEach(() => {
+    mockWSServer = createMockWSServer()
   })
 
-  afterAll(() => {
-    server.close()
+  afterEach(() => {
+    mockWSServer.server.close()
   })
 
   it('should close the WebSocket when it is not already open', () => {
     // Arrange
-    const wsgo = WSGO(`ws://localhost:${port}`)
+    const wsgo = WSGO(`ws://localhost:${mockWSServer.port}`)
 
     // Act
     wsgo.close()

test/create.test.ts

@@ -1,26 +1,21 @@
-import { describe, it, expect, beforeAll, afterAll } from 'vitest'
+import { describe, it, expect, beforeEach, afterEach } from 'vitest'
 import WSGO from '../src/index'
-import ws from 'ws'
 import { createMockWSServer } from './utils'
 
 describe('create', () => {
-  let port: number = 0
-  let server: ws.Server
+  let mockWSServer: ReturnType<typeof createMockWSServer>
 
-  beforeAll(() => {
-    const mockWSServer = createMockWSServer(port)
-
-    server = mockWSServer.server
-    port = mockWSServer.port
+  beforeEach(() => {
+    mockWSServer = createMockWSServer()
   })
 
-  afterAll(() => {
-    server.close()
+  afterEach(() => {
+    mockWSServer.server.close()
   })
 
   it('should create a WebSocket, and connect to the server when immediate = default value of', () => {
     // Act
-    const wsgo = WSGO(`ws://localhost:${port}`)
+    const wsgo = WSGO(`ws://localhost:${mockWSServer.port}`)
 
     // Assert
     expect(wsgo.ws).toBeInstanceOf(window.WebSocket)
@@ -29,7 +24,7 @@ describe('create', () => {
 
   it('should create a WebSocket, and connect to the server when immediate = true', () => {
     // Act
-    const wsgo = WSGO(`ws://localhost:${port}`, {
+    const wsgo = WSGO(`ws://localhost:${mockWSServer.port}`, {
       immediate: true,
     })
 
@@ -40,7 +35,7 @@ describe('create', () => {
 
   it('should not create WebSocket when immediate = false', () => {
     // Act
-    const wsgo = WSGO(`ws://localhost:${port}`, {
+    const wsgo = WSGO(`ws://localhost:${mockWSServer.port}`, {
       immediate: false,
     })
 

test/heartbeat.test.ts

@@ -1,23 +1,17 @@
-import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest'
+import { describe, it, expect, vi, afterEach, beforeEach } from 'vitest'
 import WSGO from '../src/index'
-import ws from 'ws'
 import { createMockWSServer } from './utils'
 
 describe('heartbeat', () => {
   const date = new Date(2000, 1, 1)
+  let mockWSServer: ReturnType<typeof createMockWSServer>
 
-  let port: number = 0
-  let server: ws.Server
-
-  beforeAll(() => {
-    const mockWSServer = createMockWSServer(port)
-
-    server = mockWSServer.server
-    port = mockWSServer.port
+  beforeEach(() => {
+    mockWSServer = createMockWSServer()
   })
 
-  afterAll(() => {
-    server.close()
+  afterEach(() => {
+    mockWSServer.server.close()
   })
 
   it('should send a ping event and receive a pong response', async () => {
@@ -26,7 +20,7 @@ describe('heartbeat', () => {
     let event: any
 
     // Arrange
-    const wsgo = WSGO(`ws://localhost:${port}`)
+    const wsgo = WSGO(`ws://localhost:${mockWSServer.port}`)
     await vi.waitFor(() => {
       vi.setSystemTime(date)
       if (wsgo.ws?.readyState !== window.WebSocket.OPEN) {
@@ -53,5 +47,30 @@ describe('heartbeat', () => {
     expect(event).toStrictEqual({ event: eventName, timeSended: Date.now(), timeReceived: Date.now() })
   })
 
-  it.todo('must close the connection if no response is received from the server')
+  it('must close the connection if no response is received from the server', async () => {
+    // Arrange
+    const wsgo = WSGO(`ws://localhost:${mockWSServer.port}`)
+    await vi.waitFor(() => {
+      if (wsgo.ws?.readyState !== window.WebSocket.OPEN) {
+        throw new Error()
+      }
+    })
+
+    // Act
+    mockWSServer.turnHeartbeat()
+    await vi.waitFor(
+      () => {
+        if (wsgo.ws?.readyState !== window.WebSocket.CLOSED) {
+          throw new Error()
+        }
+      },
+      {
+        timeout: 5000,
+        interval: 250,
+      },
+    )
+
+    // Assert
+    expect(wsgo.ws?.readyState).toBe(window.WebSocket.CLOSED)
+  })
 })