diff --git a/.changeset/bright-sockets-flow.md b/.changeset/bright-sockets-flow.md new file mode 100644 index 0000000..eb10763 --- /dev/null +++ b/.changeset/bright-sockets-flow.md @@ -0,0 +1,5 @@ +--- +"partysocket": patch +--- + +Add `enabled` prop to `usePartySocket` and `useWebSocket` hooks for conditional connection control diff --git a/packages/partysocket/src/react.ts b/packages/partysocket/src/react.ts index 0646814..1277285 100644 --- a/packages/partysocket/src/react.ts +++ b/packages/partysocket/src/react.ts @@ -7,9 +7,11 @@ import { import type { PartySocketOptions } from "./index"; import type { EventHandlerOptions } from "./use-handlers"; +import type { SocketOptions } from "./use-socket"; type UsePartySocketOptions = Omit & - EventHandlerOptions & { + EventHandlerOptions & + Pick & { host?: string | undefined; }; diff --git a/packages/partysocket/src/tests/react-hooks.test.tsx b/packages/partysocket/src/tests/react-hooks.test.tsx index 572fee3..24bd192 100644 --- a/packages/partysocket/src/tests/react-hooks.test.tsx +++ b/packages/partysocket/src/tests/react-hooks.test.tsx @@ -11,7 +11,7 @@ import usePartySocket, { useWebSocket } from "../react"; const PORT = 50128; // const URL = `ws://localhost:${PORT}`; -describe.skip("usePartySocket", () => { +describe("usePartySocket", () => { let wss: WebSocketServer; beforeAll(() => { @@ -653,9 +653,145 @@ describe.skip("usePartySocket", () => { result.current.close(); }); + + test("does not connect when enabled is false", () => { + const { result } = renderHook(() => + usePartySocket({ + host: `localhost:${PORT}`, + room: "test-room", + enabled: false + }) + ); + + expect(result.current).toBeDefined(); + expect(result.current.readyState).toBe(WebSocket.CLOSED); + }); + + test("connects when enabled is true (default)", async () => { + wss.once("connection", (ws) => { + ws.close(); + }); + + const { result } = renderHook(() => + usePartySocket({ + host: `localhost:${PORT}`, + room: "test-room", + enabled: true + }) + ); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + result.current.close(); + }); + + test("disconnects when enabled changes from true to false", async () => { + wss.once("connection", (ws) => { + // Keep connection open + }); + + const { result, rerender } = renderHook( + ({ enabled }) => + usePartySocket({ + host: `localhost:${PORT}`, + room: "test-room", + enabled + }), + { initialProps: { enabled: true } } + ); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + rerender({ enabled: false }); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.CLOSED); + }, + { timeout: 3000 } + ); + }); + + test("reconnects when enabled changes from false to true", async () => { + wss.once("connection", (ws) => { + // Keep connection open + }); + + const { result, rerender } = renderHook( + ({ enabled }) => + usePartySocket({ + host: `localhost:${PORT}`, + room: "test-room", + enabled + }), + { initialProps: { enabled: false } } + ); + + expect(result.current.readyState).toBe(WebSocket.CLOSED); + + rerender({ enabled: true }); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + result.current.close(); + }); + + test("keeps the same socket instance when enabled toggles", async () => { + wss.once("connection", () => { + // Keep connection open + }); + + const { result, rerender } = renderHook( + ({ enabled }) => + usePartySocket({ + host: `localhost:${PORT}`, + room: "test-room", + enabled + }), + { initialProps: { enabled: true } } + ); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + const socketInstance = result.current; + + rerender({ enabled: false }); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.CLOSED); + }, + { timeout: 3000 } + ); + + // Same socket instance should be reused + expect(result.current).toBe(socketInstance); + + result.current.close(); + }); }); -describe.skip("useWebSocket", () => { +describe("useWebSocket", () => { let wss: WebSocketServer; beforeAll(() => { @@ -835,4 +971,130 @@ describe.skip("useWebSocket", () => { result.current.close(); }); + + test("does not connect when enabled is false", () => { + const { result } = renderHook(() => + useWebSocket(`ws://localhost:${PORT + 1}`, undefined, { + enabled: false + }) + ); + + expect(result.current).toBeDefined(); + expect(result.current.readyState).toBe(WebSocket.CLOSED); + }); + + test("connects when enabled is true (default)", async () => { + wss.once("connection", (ws) => { + ws.close(); + }); + + const { result } = renderHook(() => + useWebSocket(`ws://localhost:${PORT + 1}`, undefined, { + enabled: true + }) + ); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + result.current.close(); + }); + + test("disconnects when enabled changes from true to false", async () => { + wss.once("connection", () => { + // Keep connection open + }); + + const { result, rerender } = renderHook( + ({ enabled }) => + useWebSocket(`ws://localhost:${PORT + 1}`, undefined, { + enabled + }), + { initialProps: { enabled: true } } + ); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + rerender({ enabled: false }); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.CLOSED); + }, + { timeout: 3000 } + ); + }); + + test("reconnects when enabled changes from false to true", async () => { + wss.once("connection", () => { + // Keep connection open + }); + + const { result, rerender } = renderHook( + ({ enabled }) => + useWebSocket(`ws://localhost:${PORT + 1}`, undefined, { + enabled + }), + { initialProps: { enabled: false } } + ); + + expect(result.current.readyState).toBe(WebSocket.CLOSED); + + rerender({ enabled: true }); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + result.current.close(); + }); + + test("keeps the same socket instance when enabled toggles", async () => { + wss.once("connection", () => { + // Keep connection open + }); + + const { result, rerender } = renderHook( + ({ enabled }) => + useWebSocket(`ws://localhost:${PORT + 1}`, undefined, { + enabled + }), + { initialProps: { enabled: true } } + ); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.OPEN); + }, + { timeout: 3000 } + ); + + const socketInstance = result.current; + + rerender({ enabled: false }); + + await waitFor( + () => { + expect(result.current.readyState).toBe(WebSocket.CLOSED); + }, + { timeout: 3000 } + ); + + // Same socket instance should be reused + expect(result.current).toBe(socketInstance); + + result.current.close(); + }); }); diff --git a/packages/partysocket/src/use-socket.ts b/packages/partysocket/src/use-socket.ts index 0cc1f93..065f00e 100644 --- a/packages/partysocket/src/use-socket.ts +++ b/packages/partysocket/src/use-socket.ts @@ -3,10 +3,16 @@ import { useEffect, useMemo, useRef, useState } from "react"; import type WebSocket from "./ws"; import type { Options } from "./ws"; +export type SocketOptions = Options & { + /** Whether the socket should be connected. Defaults to true. */ + enabled?: boolean; +}; + /** When any of the option values are changed, we should reinitialize the socket */ export const getOptionsThatShouldCauseRestartWhenChanged = ( - options: Options + options: SocketOptions ) => [ + // Note: enabled is handled separately to avoid creating a new socket on toggle options.startClosed, options.minUptime, options.maxRetries, @@ -22,7 +28,10 @@ export const getOptionsThatShouldCauseRestartWhenChanged = ( * Initializes a PartySocket (or WebSocket) and keeps it stable across renders, * but reconnects and updates the reference when any of the connection args change. */ -export function useStableSocket({ +export function useStableSocket< + T extends WebSocket, + TOpts extends SocketOptions +>({ options, createSocket, createSocketMemoKey: createOptionsMemoKey @@ -31,6 +40,9 @@ export function useStableSocket({ createSocket: (options: TOpts) => T; createSocketMemoKey: (options: TOpts) => string; }) { + // extract enabled with default value of true + const { enabled = true } = options; + // ensure we only reconnect when necessary const shouldReconnect = createOptionsMemoKey(options); const socketOptions = useMemo(() => { @@ -50,8 +62,27 @@ export function useStableSocket({ const createSocketRef = useRef(createSocket); createSocketRef.current = createSocket; + // track the previous enabled state to detect changes + const prevEnabledRef = useRef(enabled); + // finally, initialize the socket useEffect(() => { + // if disabled, close the socket and don't proceed with connection logic + if (!enabled) { + socket.close(); + prevEnabledRef.current = enabled; + return; + } + + // if enabled just changed from false to true, reconnect + if (!prevEnabledRef.current && enabled) { + socket.reconnect(); + prevEnabledRef.current = enabled; + return; + } + + prevEnabledRef.current = enabled; + // we haven't yet restarted the socket if (socketInitializedRef.current === socket) { // create new socket @@ -76,7 +107,7 @@ export function useStableSocket({ socket.close(); }; } - }, [socket, socketOptions]); + }, [socket, socketOptions, enabled]); return socket; } diff --git a/packages/partysocket/src/use-ws.ts b/packages/partysocket/src/use-ws.ts index d7fa27f..6e77fae 100644 --- a/packages/partysocket/src/use-ws.ts +++ b/packages/partysocket/src/use-ws.ts @@ -6,9 +6,10 @@ import { import WebSocket from "./ws"; import type { EventHandlerOptions } from "./use-handlers"; -import type { Options, ProtocolsProvider, UrlProvider } from "./ws"; +import type { SocketOptions } from "./use-socket"; +import type { ProtocolsProvider, UrlProvider } from "./ws"; -type UseWebSocketOptions = Options & EventHandlerOptions; +type UseWebSocketOptions = SocketOptions & EventHandlerOptions; // A React hook that wraps PartySocket export default function useWebSocket(