1
import { act, renderHook } from "@testing-library/react";
2
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
3
import { useWebSocket } from "@/client/hooks/useWebSocket";
4
import { MockWebSocket } from "../../setup/client";
5
 
6
const { createWebSocketTicketMock } = vi.hoisted(() => ({
7
  createWebSocketTicketMock: vi.fn(),
8
}));
9
 
10
vi.mock("@/client/lib/api", () => ({
11
  createWebSocketTicket: createWebSocketTicketMock,
12
}));
13
 
14
class TrackingWebSocket extends MockWebSocket {
15
  static instances: TrackingWebSocket[] = [];
16
 
17
  constructor(url: string | URL, protocols?: string | string[]) {
18
    super(url, protocols);
19
    TrackingWebSocket.instances.push(this);
20
  }
21
}
22
 
23
function decoder<T>() {
24
  return {
25
    assertDecode: vi.fn((value: unknown) => value as T),
26
  };
27
}
28
 
29
async function flushAsyncWork() {
30
  await act(async () => {
31
    await Promise.resolve();
32
    await Promise.resolve();
33
    await Promise.resolve();
34
  });
35
}
36
 
37
describe("useWebSocket", () => {
38
  beforeEach(() => {
39
    vi.restoreAllMocks();
40
    vi.useRealTimers();
41
    createWebSocketTicketMock.mockReset();
42
    TrackingWebSocket.instances = [];
43
    vi.stubGlobal("WebSocket", TrackingWebSocket);
44
  });
45
 
46
  afterEach(() => {
47
    vi.useRealTimers();
48
  });
49
 
50
  it("stays idle when disabled", () => {
51
    const { result } = renderHook(() =>
52
      useWebSocket({
53
        address: "reader@mail.test",
54
        auth: { mode: "user" as const, token: "tok_user" },
55
        enabled: false,
56
        messageCodec: decoder<{ ok: true }>(),
57
      }),
58
    );
59
 
60
    expect(result.current).toBe("idle");
61
    expect(createWebSocketTicketMock).not.toHaveBeenCalled();
62
  });
63
 
64
  it("requests a websocket ticket before connecting", async () => {
65
    createWebSocketTicketMock.mockResolvedValueOnce({
66
      ticket: "wst_123",
67
    });
68
 
69
    renderHook(() =>
70
      useWebSocket({
71
        address: "reader@mail.test",
72
        auth: { mode: "user" as const, token: "tok_user" },
73
        enabled: true,
74
        messageCodec: decoder<{ ok: true }>(),
75
      }),
76
    );
77
 
78
    await flushAsyncWork();
79
 
80
    expect(createWebSocketTicketMock).toHaveBeenCalledWith("reader@mail.test", { mode: "user", token: "tok_user" });
81
    expect(TrackingWebSocket.instances).toHaveLength(1);
82
    expect(TrackingWebSocket.instances[0]?.url).toContain("/ws?address=reader%40mail.test&ticket=wst_123");
83
  });
84
 
85
  it("transitions to open and sends heartbeat pings", async () => {
86
    createWebSocketTicketMock.mockResolvedValueOnce({
87
      ticket: "wst_123",
88
    });
89
    const onMessage = vi.fn<(message: { type: string }) => void>();
90
    const messageCodec = decoder<{ type: string }>();
91
    const { result } = renderHook(() =>
92
      useWebSocket({
93
        address: "reader@mail.test",
94
        auth: { mode: "user" as const, token: "tok_user" },
95
        enabled: true,
96
        messageCodec,
97
        onMessage,
98
      }),
99
    );
100
 
101
    await flushAsyncWork();
102
 
103
    vi.useFakeTimers();
104
    act(() => {
105
      TrackingWebSocket.instances[0]?.open();
106
    });
107
 
108
    expect(result.current).toBe("open");
109
 
110
    act(() => {
111
      TrackingWebSocket.instances[0]?.receive(JSON.stringify({ type: "new_email" }));
112
    });
113
 
114
    expect(onMessage).toHaveBeenCalledWith({ type: "new_email" });
115
 
116
    act(() => {
117
      vi.advanceTimersByTime(25_000);
118
    });
119
 
120
    expect(TrackingWebSocket.instances[0]?.send).toHaveBeenCalledWith("ping");
121
  });
122
 
123
  it("ignores malformed message payloads", async () => {
124
    createWebSocketTicketMock.mockResolvedValueOnce({
125
      ticket: "wst_123",
126
    });
127
    const onMessage = vi.fn();
128
    const messageCodec = {
129
      assertDecode: vi.fn(() => {
130
        throw new Error("invalid payload");
131
      }),
132
    };
133
 
134
    renderHook(() =>
135
      useWebSocket({
136
        address: "reader@mail.test",
137
        auth: { mode: "user" as const, token: "tok_user" },
138
        enabled: true,
139
        messageCodec,
140
        onMessage,
141
      }),
142
    );
143
 
144
    await flushAsyncWork();
145
 
146
    act(() => {
147
      TrackingWebSocket.instances[0]?.receive('{"broken":true}');
148
    });
149
 
150
    expect(onMessage).not.toHaveBeenCalled();
151
  });
152
 
153
  it("schedules reconnect after close", async () => {
154
    createWebSocketTicketMock
155
      .mockResolvedValueOnce({ ticket: "wst_first" })
156
      .mockResolvedValueOnce({ ticket: "wst_second" });
157
    vi.useFakeTimers();
158
    const setTimeoutSpy = vi.spyOn(window, "setTimeout");
159
 
160
    renderHook(() =>
161
      useWebSocket({
162
        address: "reader@mail.test",
163
        auth: { mode: "user" as const, token: "tok_user" },
164
        enabled: true,
165
        messageCodec: decoder<{ ok: true }>(),
166
      }),
167
    );
168
 
169
    await flushAsyncWork();
170
 
171
    act(() => {
172
      TrackingWebSocket.instances[0]?.dispatchEvent(new Event("close"));
173
    });
174
 
175
    expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 1000);
176
  });
177
});