From ee811844d07365ca98721be90c4e1c2c1d8623b9 Mon Sep 17 00:00:00 2001
From: sanduluca <56228534+sanduluca@users.noreply.github.com>
Date: Mon, 1 Jul 2024 18:54:32 +0300
Subject: [PATCH] feat: allow to pass custom timer for keepalive manager
 (#1896)

* feat: allow to pass custom timer for keepalive manager

* docs(readme): update `timerVariant` description

* feat: export `Timer` type

* feat(test): add test for KeepaliveManager to use provided Timer object

* fix(type): add Timer type to timerVariant on client options

---------

Co-authored-by: Sandu Luca <sandulucawork@gmail.com>
---
 README.md                   |  9 ++++++++-
 src/lib/KeepaliveManager.ts |  9 +++++++--
 src/lib/client.ts           |  4 +++-
 src/mqtt.ts                 |  1 +
 test/keepaliveManager.ts    | 20 +++++++++++++++++++-
 5 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 74a43cffb..f13a9a98e 100644
--- a/README.md
+++ b/README.md
@@ -466,7 +466,14 @@ The arguments are:
   - `messageIdProvider`: custom messageId provider. when `new UniqueMessageIdProvider()` is set, then non conflict messageId is provided.
   - `log`: custom log function. Default uses [debug](https://www.npmjs.com/package/debug) package.
   - `manualConnect`: prevents the constructor to call `connect`. In this case after the `mqtt.connect` is called you should call `client.connect` manually.
-  - `timerVariant`: defaults to `auto`, which tries to determine which timer is most appropriate for you environment, if you're having detection issues, you can set it to `worker` or `native`
+  - `timerVariant`: defaults to `auto`, which tries to determine which timer is most appropriate for you environment, if you're having detection issues, you can set it to `worker` or `native`. If none suits you, you can pass a timer object with set and clear properties:
+  ```js 
+  timerVariant: { 
+    set: (func, timer) => setInterval(func, timer),
+    clear: (id) => clearInterval(id)
+  }
+  ```
+
   - `unixSocket`: if you want to connect to a unix socket, set this to true
 
 In case mqtts (mqtt over tls) is required, the `options` object is passed through to [`tls.connect()`](http://nodejs.org/api/tls.html#tls_tls_connect_options_callback). If using a **self-signed certificate**, set `rejectUnauthorized: false`. However, be cautious as this exposes you to potential man in the middle attacks and isn't recommended for production.
diff --git a/src/lib/KeepaliveManager.ts b/src/lib/KeepaliveManager.ts
index 14daf1eb3..c68be773e 100644
--- a/src/lib/KeepaliveManager.ts
+++ b/src/lib/KeepaliveManager.ts
@@ -33,9 +33,14 @@ export default class KeepaliveManager {
 		return this._keepalive
 	}
 
-	constructor(client: MqttClient, variant: TimerVariant) {
+	constructor(client: MqttClient, variant: TimerVariant | Timer) {
 		this.client = client
-		this.timer = getTimer(variant)
+		this.timer =
+			typeof variant === 'object' &&
+			'set' in variant &&
+			'clear' in variant
+				? variant
+				: getTimer(variant)
 		this.setKeepalive(client.options.keepalive)
 	}
 
diff --git a/src/lib/client.ts b/src/lib/client.ts
index 8ea645f61..20fdce334 100644
--- a/src/lib/client.ts
+++ b/src/lib/client.ts
@@ -41,6 +41,7 @@ import TopicAliasSend from './topic-alias-send'
 import { TypedEventEmitter } from './TypedEmitter'
 import KeepaliveManager from './KeepaliveManager'
 import isBrowser, { isWebWorker } from './is-browser'
+import { Timer } from './get-timer'
 
 const setImmediate =
 	globalThis.setImmediate ||
@@ -278,8 +279,9 @@ export interface IClientOptions extends ISecureClientOptions {
 	properties?: IConnectPacket['properties']
 	/**
 	 * @description 'auto', set to 'native' or 'worker' if you're having issues with 'auto' detection
+	 * or pass a custom timer object
 	 */
-	timerVariant?: TimerVariant
+	timerVariant?: TimerVariant | Timer
 }
 
 export interface IClientPublishOptions {
diff --git a/src/mqtt.ts b/src/mqtt.ts
index e3a79abeb..6785affc9 100644
--- a/src/mqtt.ts
+++ b/src/mqtt.ts
@@ -25,3 +25,4 @@ export {
 export * from './lib/client'
 export * from './lib/shared'
 export { ReasonCodes } from './lib/handlers/ack'
+export type { Timer } from './lib/get-timer'
diff --git a/test/keepaliveManager.ts b/test/keepaliveManager.ts
index 60084f139..91d5a24a4 100644
--- a/test/keepaliveManager.ts
+++ b/test/keepaliveManager.ts
@@ -1,7 +1,7 @@
 import { afterEach, beforeEach, describe, it } from 'node:test'
 import KeepaliveManager from '../src/lib/KeepaliveManager'
 import { assert } from 'chai'
-import { useFakeTimers, spy, mock } from 'sinon'
+import { useFakeTimers, spy, stub } from 'sinon'
 import { MqttClient } from 'src'
 
 function mockedClient(keepalive: number) {
@@ -105,4 +105,22 @@ describe('KeepaliveManager', () => {
 		assert.equal(manager.keepalive, 10000)
 		assert.equal(manager.intervalEvery, 5000)
 	})
+
+	it('should use provided Timer object', () => {
+		const keepalive = 10 // seconds
+		const customTimer = {
+			set: stub().returns(123),
+			clear: stub(),
+		}
+		const manager = new KeepaliveManager(
+			mockedClient(keepalive),
+			customTimer,
+		)
+		assert.equal(manager['timer'], customTimer)
+		assert.equal(customTimer.set.callCount, 1)
+		assert.equal(manager['timerId'], 123)
+
+		manager.destroy()
+		assert.equal(customTimer.clear.called, true)
+	})
 })