From 0c35ae318ef20ee69c8e3ea6117b3472c64e6683 Mon Sep 17 00:00:00 2001
From: Davi de Medeiros <davi.demedeiros@here.com>
Date: Mon, 13 Dec 2021 22:26:01 +0100
Subject: [PATCH 1/2] IMPROVEMENT: test the onClose callback for each hook
 consuming component is invoked when all subscribers disconnect

---
 src/lib/use-websocket.test.ts | 41 +++++++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/lib/use-websocket.test.ts b/src/lib/use-websocket.test.ts
index 4180b53..42e7c79 100644
--- a/src/lib/use-websocket.test.ts
+++ b/src/lib/use-websocket.test.ts
@@ -272,7 +272,44 @@ test('shared websockets each have callbacks invoked as if unshared', async (done
   done();
 })
 
-test('Options#fromSocketIO changes the WS url to support socket.io\'s required query params', async (done) => {
+
+test('shared websockets have their onClose callbacks invoked when all subscriber components disconnect', async (done) => {
+  const initialProps = { initialValue: true };
+  const onCloseFn1 = jest.fn();
+  const onCloseFn2 = jest.fn();
+  const onCloseFn3 = jest.fn();
+
+  const { rerender: rerender1 } = renderHook(
+    ({ initialValue }) => useWebSocket(URL, { share: true, onClose: onCloseFn1 }, initialValue),
+    { initialProps }
+  );
+  await server.connected;
+
+  const { rerender: rerender2 } = renderHook(
+    ({ initialValue }) => useWebSocket(URL, { share: true, onClose: onCloseFn2 }, initialValue),
+    { initialProps }
+  );
+  await server.connected;
+
+  const { rerender: rerender3 } = renderHook(
+    ({ initialValue }) => useWebSocket(URL, { share: true, onClose: onCloseFn3 }, initialValue),
+    { initialProps }
+  );
+  await server.connected;
+
+  rerender1({ initialValue: false });
+  rerender2({ initialValue: false });
+  rerender3({ initialValue: false });
+
+  await sleep(500);
+
+  expect(onCloseFn1).toHaveBeenCalledTimes(1);
+  expect(onCloseFn2).toHaveBeenCalledTimes(1);
+  expect(onCloseFn3).toHaveBeenCalledTimes(1);
+  done();
+});
+
+test("Options#fromSocketIO changes the WS url to support socket.io's required query params", async (done) => {
   options.fromSocketIO = true;
 
   const {
@@ -500,4 +537,4 @@ test('Options#eventSourceOptions, if provided, instantiates an EventSource inste
   done();
 });
 
-//TODO: Write companion tests for useSocketIO
\ No newline at end of file
+//TODO: Write companion tests for useSocketIO

From 3334944032b9cba3e1342cf1c394026c2bf694bc Mon Sep 17 00:00:00 2001
From: Davi de Medeiros <davi.demedeiros@here.com>
Date: Mon, 13 Dec 2021 19:55:47 +0100
Subject: [PATCH 2/2] IMPROVEMENT: register socket's close event handler
 earlier instead of on last unsubscription

also use addEventListener instead of reassigning and losing the previous socket-like close handlers

this was causing a bug where only the first subscriber of a shared socket would have its handler invoked,
since its listener was only being registered when cleaning unsubscribers for the last subscription for a url
---
 src/lib/create-or-join.ts | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/lib/create-or-join.ts b/src/lib/create-or-join.ts
index ed9ac88..2bf4218 100644
--- a/src/lib/create-or-join.ts
+++ b/src/lib/create-or-join.ts
@@ -17,17 +17,21 @@ const cleanSubscribers = (
 ) => {
   return () => {
     removeSubscriber(url, subscriber);
+
+    const socketLike = sharedWebSockets[url];
+
+    if (socketLike instanceof WebSocket) {
+      socketLike.addEventListener('close', (event: WebSocketEventMap['close']) => {
+        if (optionsRef.current.onClose) {
+          optionsRef.current.onClose(event);
+        }
+
+        subscriber.setReadyState(ReadyState.CLOSED);
+      });
+    }
+
     if (!hasSubscribers(url)) {
       try {
-        const socketLike = sharedWebSockets[url];
-        if (socketLike instanceof WebSocket) {
-          socketLike.onclose = (event: WebSocketEventMap['close']) => {
-            if (optionsRef.current.onClose) {
-              optionsRef.current.onClose(event);
-            }
-            setReadyState(ReadyState.CLOSED);
-          };
-        }
         socketLike.close();
       } catch (e) {
 
@@ -83,7 +87,7 @@ export const createOrJoinSocket = (
       reconnectCount,
       reconnect: startRef,
     };
-  
+
     addSubscriber(url, subscriber);
 
     return cleanSubscribers(