fix useWebSocket error (#2176) (#2177)

* fix: useWebSocket don't call connect when socketUrl is empty (#2176)

* fix: useWebSocket unlimited reconnect when update socketUrl (#2176)

* fix: disconnect not works

* docs: modify useWebSocket docs

Author: YinDongFang

packages/hooks/src/useWebSocket/__tests__/index.test.ts

@@ -7,7 +7,7 @@ const promise: Promise<void> = new Promise((resolve) => resolve());
 const wsUrl = 'ws://localhost:9999';
 
 describe('useWebSocket', () => {
-  afterAll(() => {
+  afterEach(() => {
     WS.clean();
   });
 
@@ -44,6 +44,21 @@ describe('useWebSocket', () => {
     expect(hooks.result.current.readyState).toBe(ReadyState.Closed);
   });
 
+  it('disconnect should work', async () => {
+    const wsServer = new WS(wsUrl);
+    const hooks = renderHook(() => useWebSocket(wsUrl));
+
+    // connect
+    expect(hooks.result.current.readyState).toBe(ReadyState.Connecting);
+    await act(() => wsServer.connected);
+    expect(hooks.result.current.readyState).toBe(ReadyState.Open);
+
+    // disconnect
+    act(() => hooks.result.current.disconnect());
+    await act(() => wsServer.closed);
+    expect(hooks.result.current.readyState).toBe(ReadyState.Closed);
+  });
+
   it('useWebSocket should be manually triggered', async () => {
     const wsServer = new WS(wsUrl);
 
@@ -68,4 +83,65 @@ describe('useWebSocket', () => {
 
     act(() => wsServer.close());
   });
+
+  it('should not call connect when initial socketUrl is empty', async () => {
+    const wsServer = new WS(wsUrl);
+    const onOpen = jest.fn();
+    const onClose = jest.fn();
+
+    let url = '';
+    const hooks = renderHook(() => useWebSocket(url, { onOpen, onClose }));
+
+    await act(async () => {
+      await sleep(1000);
+    });
+
+    expect(hooks.result.current.readyState).toBe(ReadyState.Closed);
+
+    url = wsUrl;
+    hooks.rerender();
+
+    await act(async () => {
+      await wsServer.connected;
+    });
+
+    expect(hooks.result.current.readyState).toBe(ReadyState.Open);
+    expect(onOpen).toBeCalledTimes(1);
+
+    act(() => wsServer.close());
+  });
+
+  it('change socketUrl should connect correctly', async () => {
+    const wsUrl1 = 'ws://localhost:8888';
+    const wsServer1 = new WS(wsUrl);
+    const wsServer2 = new WS(wsUrl1);
+
+    const onOpen = jest.fn();
+    const onClose = jest.fn();
+
+    let url = wsUrl;
+    const hooks = renderHook(() => useWebSocket(url, { onOpen, onClose, reconnectInterval: 300 }));
+
+    expect(hooks.result.current.readyState).toBe(ReadyState.Connecting);
+    await act(async () => {
+      await wsServer1.connected;
+    });
+    expect(hooks.result.current.readyState).toBe(ReadyState.Open);
+
+    url = wsUrl1;
+    hooks.rerender();
+    await act(async () => {
+      await wsServer2.connected;
+    });
+    expect(hooks.result.current.readyState).toBe(ReadyState.Open);
+
+    await act(async () => {
+      await sleep(3000);
+    });
+    expect(onOpen).toBeCalledTimes(2);
+    expect(onClose).toBeCalledTimes(1);
+
+    act(() => wsServer1.close());
+    act(() => wsServer2.close());
+  });
 });

packages/hooks/src/useWebSocket/index.en-US.md

@@ -35,9 +35,9 @@ interface Options {
 
 interface Result {
   latestMessage?: WebSocketEventMap['message'];
-  sendMessage?: WebSocket['send'];
-  disconnect?: () => void;
-  connect?: () => void;
+  sendMessage: WebSocket['send'];
+  disconnect: () => void;
+  connect: () => void;
   readyState: ReadyState;
   webSocketIns?: WebSocket;
 }

packages/hooks/src/useWebSocket/index.ts

@@ -24,9 +24,9 @@ export interface Options {
 
 export interface Result {
   latestMessage?: WebSocketEventMap['message'];
-  sendMessage?: WebSocket['send'];
-  disconnect?: () => void;
-  connect?: () => void;
+  sendMessage: WebSocket['send'];
+  disconnect: () => void;
+  connect: () => void;
   readyState: ReadyState;
   webSocketIns?: WebSocket;
 }
@@ -52,8 +52,6 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
   const reconnectTimerRef = useRef<ReturnType<typeof setTimeout>>();
   const websocketRef = useRef<WebSocket>();
 
-  const unmountedRef = useRef(false);
-
   const [latestMessage, setLatestMessage] = useState<WebSocketEventMap['message']>();
   const [readyState, setReadyState] = useState<ReadyState>(ReadyState.Closed);
 
@@ -87,35 +85,38 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
     setReadyState(ReadyState.Connecting);
 
     ws.onerror = (event) => {
-      if (unmountedRef.current) {
+      if (websocketRef.current !== ws) {
         return;
       }
       reconnect();
       onErrorRef.current?.(event, ws);
       setReadyState(ws.readyState || ReadyState.Closed);
     };
     ws.onopen = (event) => {
-      if (unmountedRef.current) {
+      if (websocketRef.current !== ws) {
         return;
       }
       onOpenRef.current?.(event, ws);
       reconnectTimesRef.current = 0;
       setReadyState(ws.readyState || ReadyState.Open);
     };
     ws.onmessage = (message: WebSocketEventMap['message']) => {
-      if (unmountedRef.current) {
+      if (websocketRef.current !== ws) {
         return;
       }
       onMessageRef.current?.(message, ws);
       setLatestMessage(message);
     };
     ws.onclose = (event) => {
-      if (unmountedRef.current) {
-        return;
-      }
-      reconnect();
       onCloseRef.current?.(event, ws);
-      setReadyState(ws.readyState || ReadyState.Closed);
+      // closed by server
+      if (websocketRef.current === ws) {
+        reconnect();
+      }
+      // closed by disconnect or closed by server
+      if (!websocketRef.current || websocketRef.current === ws) {
+        setReadyState(ws.readyState || ReadyState.Closed);
+      }
     };
 
     websocketRef.current = ws;
@@ -141,16 +142,16 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
 
     reconnectTimesRef.current = reconnectLimit;
     websocketRef.current?.close();
+    websocketRef.current = undefined;
   };
 
   useEffect(() => {
-    if (!manual) {
+    if (!manual && socketUrl) {
       connect();
     }
   }, [socketUrl, manual]);
 
   useUnmount(() => {
-    unmountedRef.current = true;
     disconnect();
   });
 

packages/hooks/src/useWebSocket/index.zh-CN.md

@@ -35,9 +35,9 @@ interface Options {
 
 interface Result {
   latestMessage?: WebSocketEventMap['message'];
-  sendMessage?: WebSocket['send'];
-  disconnect?: () => void;
-  connect?: () => void;
+  sendMessage: WebSocket['send'];
+  disconnect: () => void;
+  connect: () => void;
   readyState: ReadyState;
   webSocketIns?: WebSocket;
 }