feat: update code

Author: meet-student

packages/hooks/src/useSse/__tests__/index.spec.ts

@@ -1,5 +1,6 @@
+// useSse.test.ts
 import { act, renderHook } from '@testing-library/react';
-import { afterEach, describe, expect, test, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';
 import useSse, { ReadyState } from '../index';
 
 class MockEventSource {
@@ -10,6 +11,7 @@ class MockEventSource {
   onmessage: ((this: EventSource, ev: MessageEvent) => any) | null = null;
   onerror: ((this: EventSource, ev: Event) => any) | null = null;
   private listeners: Record<string, Array<(ev: Event) => void>> = {};
+  private openTimeout?: NodeJS.Timeout;
 
   static CONNECTING = 0;
   static OPEN = 1;
@@ -19,9 +21,10 @@ class MockEventSource {
     this.url = url;
     this.withCredentials = Boolean(init?.withCredentials);
     this.readyState = MockEventSource.CONNECTING;
-    setTimeout(() => {
+
+    this.openTimeout = setTimeout(() => {
       this.readyState = MockEventSource.OPEN;
-      this.onopen && this.onopen(new Event('open'));
+      this.onopen?.(new Event('open'));
     }, 10);
   }
 
@@ -35,69 +38,153 @@ class MockEventSource {
   }
 
   emitMessage(data: any) {
-    this.onmessage && this.onmessage(new MessageEvent('message', { data }));
+    if (this.readyState !== MockEventSource.OPEN) return;
+    this.onmessage?.(new MessageEvent('message', { data }));
   }
 
   emitError() {
-    this.onerror && this.onerror(new Event('error'));
+    this.onerror?.(new Event('error'));
+  }
+
+  emitRetry(ms: number) {
+    const ev = new MessageEvent('message', { data: '' });
+    (ev as any).retry = ms;
+    this.onmessage?.(ev);
   }
 
   close() {
     this.readyState = MockEventSource.CLOSED;
+    if (this.openTimeout) clearTimeout(this.openTimeout);
   }
 }
 
-describe('useSse', () => {
+describe('useSse Hook', () => {
   const OriginalEventSource = (globalThis as any).EventSource;
 
+  beforeEach(() => {
+    vi.useFakeTimers();
+    (globalThis as any).EventSource = MockEventSource;
+  });
+
   afterEach(() => {
+    vi.runAllTimers();
+    vi.useRealTimers();
     (globalThis as any).EventSource = OriginalEventSource;
+    vi.restoreAllMocks();
   });
 
-  test('should connect and receive message', async () => {
-    (globalThis as any).EventSource = MockEventSource as any;
+  test('should connect and receive message', () => {
+    const hook = renderHook(() => useSse('/sse'));
+    expect(hook.result.current.readyState).toBe(ReadyState.Connecting);
+
+    act(() => vi.advanceTimersByTime(20));
+    expect(hook.result.current.readyState).toBe(ReadyState.Open);
+
+    act(() => {
+      const es = hook.result.current.eventSource as unknown as MockEventSource;
+      es.emitMessage('hello');
+    });
+    expect(hook.result.current.latestMessage?.data).toBe('hello');
 
-    const hooks = renderHook(() => useSse('/sse'));
+    act(() => hook.result.current.disconnect());
+    expect(hook.result.current.readyState).toBe(ReadyState.Closed);
+  });
 
-    // not manual: should start connecting immediately
-    expect(hooks.result.current.readyState).toBe(ReadyState.Connecting);
+  test('manual mode should not auto connect', () => {
+    const hook = renderHook(() => useSse('/sse', { manual: true }));
+    expect(hook.result.current.readyState).toBe(ReadyState.Closed);
 
-    await act(async () => {
-      await new Promise((r) => setTimeout(r, 20));
+    act(() => {
+      hook.result.current.connect();
+      vi.advanceTimersByTime(20);
     });
+    expect(hook.result.current.readyState).toBe(ReadyState.Open);
 
-    expect(hooks.result.current.readyState).toBe(ReadyState.Open);
+    act(() => hook.result.current.disconnect());
+  });
+
+  test('should handle custom events', () => {
+    const onEvent = vi.fn();
+    const hook = renderHook(() => useSse('/sse', { onEvent }));
+    act(() => vi.advanceTimersByTime(20));
 
     act(() => {
-      const es = hooks.result.current.eventSource as unknown as MockEventSource;
-      es.emitMessage('hello');
+      const es = hook.result.current.eventSource as unknown as MockEventSource;
+      es.dispatchEvent('custom', new MessageEvent('custom', { data: 'foo' }));
     });
-    expect(hooks.result.current.latestMessage?.data).toBe('hello');
+
+    expect(onEvent).toHaveBeenCalledWith(
+      'custom',
+      expect.objectContaining({ data: 'foo' }),
+      expect.any(MockEventSource),
+    );
+
+    act(() => hook.result.current.disconnect());
   });
 
-  test('manual should not auto connect', async () => {
-    (globalThis as any).EventSource = MockEventSource as any;
+  test('should reconnect on error respecting reconnectLimit', () => {
+    const hook = renderHook(() => useSse('/sse', { reconnectLimit: 1, reconnectInterval: 5 }));
+    act(() => vi.advanceTimersByTime(20));
+    expect(hook.result.current.readyState).toBe(ReadyState.Open);
+
+    act(() => {
+      const es = hook.result.current.eventSource as unknown as MockEventSource;
+      es.emitError();
+      vi.advanceTimersByTime(20);
+    });
+
+    expect(
+      [ReadyState.Reconnecting, ReadyState.Open].includes(hook.result.current.readyState),
+    ).toBe(true);
+
+    act(() => hook.result.current.disconnect());
+  });
 
-    const hooks = renderHook(() => useSse('/sse', { manual: true }));
-    expect(hooks.result.current.readyState).toBe(ReadyState.Closed);
+  test('should respect server retry when enabled', () => {
+    const hook = renderHook(() =>
+      useSse('/sse', { reconnectLimit: 1, reconnectInterval: 5, respectServerRetry: true }),
+    );
+    act(() => vi.advanceTimersByTime(20));
+    expect(hook.result.current.readyState).toBe(ReadyState.Open);
 
-    await act(async () => {
-      hooks.result.current.connect();
-      await new Promise((r) => setTimeout(r, 20));
+    act(() => {
+      const es = hook.result.current.eventSource as unknown as MockEventSource;
+      es.emitRetry(50);
+      es.emitError();
+      vi.advanceTimersByTime(60);
     });
 
-    expect(hooks.result.current.readyState).toBe(ReadyState.Open);
+    expect(
+      [ReadyState.Reconnecting, ReadyState.Open].includes(hook.result.current.readyState),
+    ).toBe(true);
+
+    act(() => hook.result.current.disconnect());
   });
 
-  test('disconnect should close', async () => {
-    (globalThis as any).EventSource = MockEventSource as any;
+  test('should trigger all callbacks', () => {
+    const onOpen = vi.fn();
+    const onMessage = vi.fn();
+    const onError = vi.fn();
+    const onReconnect = vi.fn();
+
+    const hook = renderHook(() => useSse('/sse', { onOpen, onMessage, onError, onReconnect }));
+    act(() => vi.advanceTimersByTime(20));
+    expect(onOpen).toHaveBeenCalled();
 
-    const hooks = renderHook(() => useSse('/sse'));
-    await act(async () => {
-      await new Promise((r) => setTimeout(r, 20));
+    act(() => {
+      const es = hook.result.current.eventSource as unknown as MockEventSource;
+      es.emitMessage('world');
     });
-    expect(hooks.result.current.readyState).toBe(ReadyState.Open);
-    act(() => hooks.result.current.disconnect());
-    expect(hooks.result.current.readyState).toBe(ReadyState.Closed);
+    expect(onMessage).toHaveBeenCalled();
+
+    act(() => {
+      const es = hook.result.current.eventSource as unknown as MockEventSource;
+      es.emitError();
+      vi.advanceTimersByTime(20);
+    });
+    expect(onError).toHaveBeenCalled();
+    expect(onReconnect).toHaveBeenCalled();
+
+    act(() => hook.result.current.disconnect());
   });
 });

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

@@ -5,50 +5,44 @@ nav:
 
 # useSse
 
-Listen to Server-Sent Events (SSE) stream with auto reconnect and lifecycle helpers.
-
-### Examples
-
-```tsx
-import React, { useMemo, useRef } from 'react';
-import { useSse } from 'ahooks';
-
-export default () => {
-  const historyRef = useRef<any[]>([]);
-  const { readyState, latestMessage, connect, disconnect } = useSse('/api/sse');
-
-  historyRef.current = useMemo(() => historyRef.current.concat(latestMessage), [latestMessage]);
-
-  return (
-    <div>
-      <button onClick={() => connect()} style={{ marginRight: 8 }}>connect</button>
-      <button onClick={() => disconnect()} style={{ marginRight: 8 }}>disconnect</button>
-      <div>readyState: {readyState}</div>
-      <div style={{ marginTop: 8 }}>
-        {historyRef.current.map((m, i) => (
-          <p key={i}>{m?.data}</p>
-        ))}
-      </div>
-    </div>
-  );
-};
-```
+A hook for Server-Sent Events (SSE), which supports automatic reconnect and message callbacks.
+
+## Examples
+
+### Basic Usage
 
-### API
+<code src="./demo/demo1.tsx" />
 
-```ts
+## API
+
+```typescript
 const { readyState, latestMessage, connect, disconnect, eventSource } = useSse(
   url: string,
-  options?: {
-    manual?: boolean;
-    withCredentials?: boolean;
-    reconnectLimit?: number;
-    reconnectInterval?: number; // ms
-    events?: string[]; // named events
-    onOpen?: (ev: Event, instance: EventSource) => void;
-    onMessage?: (msg: MessageEvent, instance: EventSource) => void;
-    onError?: (ev: Event, instance: EventSource) => void;
-    onEvent?: (eventName: string, ev: MessageEvent, instance: EventSource) => void;
-  }
+  options?: UseSseOptions
 )
 ```
+
+### Options
+
+| Property             | Description                                             | Type                                           | Default      |
+| -------------------- | ------------------------------------------------------- | ---------------------------------------------- | ------------ |
+| manual               | Whether to connect manually                             | `boolean`                                      | `false`      |
+| withCredentials      | Whether to send cross-domain requests with credentials  | `boolean`                                      | `false`      |
+| reconnectLimit       | Maximum number of reconnection attempts                 | `number`                                       | `3`          |
+| reconnectInterval    | Reconnection interval (in milliseconds)                 | `number`                                       | `3000`       |
+| respectServerRetry   | Whether to respect the retry time sent by the server    | `boolean`                                      | `false`      |
+| onOpen               | Callback when the connection is successfully established| `(es: EventSource) => void`                    | -            |
+| onMessage            | Callback when a message is received                     | `(ev: MessageEvent, es: EventSource) => void`  | -            |
+| onError              | Callback when an error occurs                           | `(ev: Event, es: EventSource) => void`         | -            |
+| onReconnect          | Callback when a reconnection occurs                     | `(attempt: number, es: EventSource) => void`   | -            |
+| onEvent              | Callback for custom events                              | `(event: string, ev: MessageEvent, es: EventSource) => void` | - |
+
+### Result
+
+| Property      | Description                                | Type                                  |
+| ------------- | ------------------------------------------ | ------------------------------------- |
+| readyState    | The current connection state               | `ReadyState` (0: Connecting, 1: Open, 2: Closed, 3: Reconnecting) |
+| latestMessage | The latest message received                | `MessageEvent` \| `null`              |
+| connect       | Function to manually connect               | `() => void`                          |
+| disconnect    | Function to manually disconnect            | `() => void`                          |
+| eventSource   | The native EventSource instance            | `EventSource` \| `null`               |