fix useWebSocket error (#2176)

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

@@ -68,4 +68,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.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,14 +52,13 @@ 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);
 
   const reconnect = () => {
     if (
       reconnectTimesRef.current < reconnectLimit &&
+      websocketRef.current &&
       websocketRef.current?.readyState !== ReadyState.Open
     ) {
       if (reconnectTimerRef.current) {
@@ -87,34 +86,34 @@ 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) {
+      onCloseRef.current?.(event, ws);
+      if (websocketRef.current !== ws) {
         return;
       }
       reconnect();
-      onCloseRef.current?.(event, ws);
       setReadyState(ws.readyState || ReadyState.Closed);
     };
 
@@ -139,18 +138,17 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
       clearTimeout(reconnectTimerRef.current);
     }
 
-    reconnectTimesRef.current = reconnectLimit;
     websocketRef.current?.close();
+    websocketRef.current = undefined;
   };
 
   useEffect(() => {
-    if (!manual) {
+    if (!manual && socketUrl) {
       connect();
     }
   }, [socketUrl, manual]);
 
   useUnmount(() => {
-    unmountedRef.current = true;
     disconnect();
   });