useWebSocket.test.ts 32 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024
  1. /**
  2. * Tests for the useWebSocket hook.
  3. *
  4. * Tests WebSocket connection management and message handling.
  5. * Uses vitest.mock to mock the entire module before MSW can intercept.
  6. */
  7. import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
  8. import { renderHook, waitFor, act, screen } from '@testing-library/react';
  9. import React from 'react';
  10. import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
  11. import { ToastProvider } from '../../contexts/ToastContext';
  12. // Track WebSocket instances created during tests
  13. let wsInstances: MockWebSocket[] = [];
  14. let originalWebSocket: typeof WebSocket;
  15. // Mock react-i18next BEFORE any modules that use it are imported
  16. vi.mock('react-i18next', () => ({
  17. useTranslation: () => ({
  18. t: (key: string, options?: Record<string, unknown>) => {
  19. if (key === 'printers.toast.missingSpoolAssignment' && options) {
  20. const { printer, slots } = options as { printer: string; slots: string };
  21. return `Missing assignments for ${printer}: ${slots}`;
  22. }
  23. if (key === 'printers.toast.killSwitchTriggered' && options) {
  24. const { printer, filename } = options as { printer: string; filename: string };
  25. return `Billing kill switch stopped ${filename} on ${printer}`;
  26. }
  27. return key;
  28. },
  29. i18n: {},
  30. }),
  31. }));
  32. // Enhanced MockWebSocket that tracks instances
  33. class MockWebSocket {
  34. static readonly CONNECTING = 0;
  35. static readonly OPEN = 1;
  36. static readonly CLOSING = 2;
  37. static readonly CLOSED = 3;
  38. readyState = MockWebSocket.CONNECTING;
  39. onopen: ((event: Event) => void) | null = null;
  40. onclose: ((event: CloseEvent) => void) | null = null;
  41. onmessage: ((event: MessageEvent) => void) | null = null;
  42. onerror: ((event: Event) => void) | null = null;
  43. url: string;
  44. constructor(url: string) {
  45. this.url = url;
  46. wsInstances.push(this);
  47. }
  48. send = vi.fn();
  49. close = vi.fn(() => {
  50. this.readyState = MockWebSocket.CLOSED;
  51. if (this.onclose) {
  52. this.onclose(new CloseEvent('close'));
  53. }
  54. });
  55. // Required by MSW's interceptor - these are no-ops but prevent the error
  56. addEventListener = vi.fn();
  57. removeEventListener = vi.fn();
  58. // Helper to simulate connection opening
  59. open() {
  60. this.readyState = MockWebSocket.OPEN;
  61. if (this.onopen) {
  62. this.onopen(new Event('open'));
  63. }
  64. }
  65. // Helper to simulate the server closing with a specific code (e.g. 4401,
  66. // the /ws auth-rejection close code).
  67. simulateClose(code: number) {
  68. this.readyState = MockWebSocket.CLOSED;
  69. if (this.onclose) {
  70. this.onclose(new CloseEvent('close', { code }));
  71. }
  72. }
  73. // Helper to simulate receiving a message
  74. simulateMessage(data: unknown) {
  75. if (this.onmessage) {
  76. this.onmessage(
  77. new MessageEvent('message', {
  78. data: JSON.stringify(data),
  79. })
  80. );
  81. }
  82. }
  83. }
  84. // Create test QueryClient
  85. function createTestQueryClient() {
  86. return new QueryClient({
  87. defaultOptions: {
  88. queries: {
  89. retry: false,
  90. gcTime: 0,
  91. },
  92. },
  93. });
  94. }
  95. // Wrapper with QueryClient and ToastProvider for hook testing
  96. function createWrapper(queryClient: QueryClient) {
  97. return function Wrapper({ children }: { children: React.ReactNode }) {
  98. return React.createElement(
  99. ToastProvider,
  100. {},
  101. React.createElement(
  102. QueryClientProvider,
  103. { client: queryClient },
  104. children
  105. )
  106. );
  107. };
  108. }
  109. /**
  110. * After GHSA-r2qv, useWebSocket awaits a ws-token fetch before constructing
  111. * the WebSocket. The MockWebSocket isn't pushed into ``wsInstances`` until
  112. * that promise resolves. ``waitFor`` from testing-library uses real-time
  113. * polling and so wedges under ``vi.useFakeTimers()``; flushing microtasks
  114. * manually works under both real and fake timers because Promise resolution
  115. * runs on the microtask queue, not on the mocked clock.
  116. *
  117. * Two iterations suffice for ``await fetch(...)`` → ``await resp.json()``;
  118. * a small headroom lets future awaits land here without changing every
  119. * call site.
  120. */
  121. async function waitForWs(): Promise<MockWebSocket> {
  122. for (let i = 0; i < 10 && wsInstances.length === 0; i++) {
  123. await Promise.resolve();
  124. }
  125. const ws = wsInstances[wsInstances.length - 1];
  126. if (!ws) {
  127. throw new Error('WebSocket was not constructed after microtask flush');
  128. }
  129. return ws;
  130. }
  131. describe('useWebSocket hook', () => {
  132. let queryClient: QueryClient;
  133. beforeEach(() => {
  134. vi.clearAllMocks();
  135. wsInstances = [];
  136. queryClient = createTestQueryClient();
  137. // Save original and install mock
  138. originalWebSocket = globalThis.WebSocket;
  139. globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket;
  140. // After GHSA-r2qv, useWebSocket fetches a ws-token via api.getWebSocketToken
  141. // before opening the socket. ``api.request`` reads ``response.headers``
  142. // and ``response.status``; the stub must expose those (a missing
  143. // ``headers`` field throws inside request() and the silent catch in
  144. // useWebSocket then proceeds with an undefined token, so the assertion
  145. // "URL contains ?token=" fails without making the cause obvious).
  146. vi.stubGlobal(
  147. 'fetch',
  148. vi.fn(async () => ({
  149. ok: true,
  150. status: 200,
  151. statusText: 'OK',
  152. headers: { get: () => null },
  153. json: async () => ({ token: 'test-ws-token' }),
  154. })),
  155. );
  156. });
  157. afterEach(() => {
  158. vi.restoreAllMocks();
  159. vi.unstubAllGlobals();
  160. // Restore original WebSocket
  161. globalThis.WebSocket = originalWebSocket;
  162. });
  163. describe('WebSocket Mock', () => {
  164. it('creates WebSocket with correct URL', () => {
  165. const ws = new MockWebSocket('ws://test.local/ws');
  166. expect(ws.url).toBe('ws://test.local/ws');
  167. });
  168. it('starts in CONNECTING state', () => {
  169. const ws = new MockWebSocket('ws://test.local/ws');
  170. expect(ws.readyState).toBe(MockWebSocket.CONNECTING);
  171. });
  172. it('transitions to OPEN state', () => {
  173. const ws = new MockWebSocket('ws://test.local/ws');
  174. const onOpen = vi.fn();
  175. ws.onopen = onOpen;
  176. ws.open();
  177. expect(ws.readyState).toBe(MockWebSocket.OPEN);
  178. expect(onOpen).toHaveBeenCalled();
  179. });
  180. it('can receive messages', () => {
  181. const ws = new MockWebSocket('ws://test.local/ws');
  182. const onMessage = vi.fn();
  183. ws.onmessage = onMessage;
  184. ws.open();
  185. ws.simulateMessage({ type: 'status', data: { connected: true } });
  186. expect(onMessage).toHaveBeenCalled();
  187. });
  188. it('can close connection', () => {
  189. const ws = new MockWebSocket('ws://test.local/ws');
  190. const onClose = vi.fn();
  191. ws.onclose = onClose;
  192. ws.close();
  193. expect(ws.readyState).toBe(MockWebSocket.CLOSED);
  194. expect(onClose).toHaveBeenCalled();
  195. });
  196. it('tracks all instances', () => {
  197. wsInstances = [];
  198. new MockWebSocket('ws://a');
  199. new MockWebSocket('ws://b');
  200. expect(wsInstances.length).toBe(2);
  201. });
  202. });
  203. describe('hook connection', () => {
  204. it('connects to WebSocket on mount', async () => {
  205. const { useWebSocket } = await import('../../hooks/useWebSocket');
  206. renderHook(() => useWebSocket(), {
  207. wrapper: createWrapper(queryClient),
  208. });
  209. const ws = await waitForWs();
  210. expect(ws).toBeDefined();
  211. expect(ws.url).toContain('/api/v1/ws');
  212. // GHSA-r2qv: the ws-token mint result is appended as ?token=...
  213. expect(ws.url).toContain('token=test-ws-token');
  214. });
  215. it('reports connected state when WebSocket opens', async () => {
  216. const { useWebSocket } = await import('../../hooks/useWebSocket');
  217. const { result } = renderHook(() => useWebSocket(), {
  218. wrapper: createWrapper(queryClient),
  219. });
  220. // Initially not connected
  221. expect(result.current.isConnected).toBe(false);
  222. // Simulate connection opening
  223. const ws = await waitForWs();
  224. act(() => {
  225. ws.open();
  226. });
  227. await waitFor(() => {
  228. expect(result.current.isConnected).toBe(true);
  229. });
  230. });
  231. });
  232. describe('message handling', () => {
  233. it('updates printer status in query cache on printer_status message', async () => {
  234. // Test the printer status update logic directly using setQueryData
  235. // The WebSocket handler with throttling is complex to test with fake timers,
  236. // so we test the core behavior directly
  237. // Simulate what the throttled update does
  238. queryClient.setQueryData(
  239. ['printerStatus', 1],
  240. (old: Record<string, unknown> | undefined) => {
  241. const statusData = { state: 'IDLE', progress: 0 };
  242. const merged = { ...old, ...statusData };
  243. return merged;
  244. }
  245. );
  246. // Check query cache was updated
  247. const cachedData = queryClient.getQueryData(['printerStatus', 1]);
  248. expect(cachedData).toEqual({ state: 'IDLE', progress: 0 });
  249. });
  250. it('preserves wifi_signal when new value is null', async () => {
  251. // Test the wifi_signal preservation logic directly on QueryClient
  252. // The throttled WebSocket handler makes this hard to test end-to-end
  253. // This tests that the merge logic correctly preserves wifi_signal
  254. // Set initial data with wifi_signal
  255. queryClient.setQueryData(['printerStatus', 1], {
  256. wifi_signal: -65,
  257. state: 'IDLE',
  258. });
  259. // Simulate what the throttled update does - use setQueryData with updater function
  260. queryClient.setQueryData(
  261. ['printerStatus', 1],
  262. (old: Record<string, unknown> | undefined) => {
  263. const statusData = { state: 'RUNNING', wifi_signal: null };
  264. const merged = { ...old, ...statusData };
  265. // This is the preservation logic from useWebSocket
  266. if (merged.wifi_signal == null && old?.wifi_signal != null) {
  267. merged.wifi_signal = old.wifi_signal;
  268. }
  269. return merged;
  270. }
  271. );
  272. const cachedData = queryClient.getQueryData(['printerStatus', 1]) as Record<
  273. string,
  274. unknown
  275. >;
  276. expect(cachedData.wifi_signal).toBe(-65); // Preserved
  277. expect(cachedData.state).toBe('RUNNING'); // Updated
  278. });
  279. it('invalidates archives on print_complete message', async () => {
  280. vi.useFakeTimers();
  281. const { useWebSocket } = await import('../../hooks/useWebSocket');
  282. const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries');
  283. renderHook(() => useWebSocket(), {
  284. wrapper: createWrapper(queryClient),
  285. });
  286. const ws = await waitForWs();
  287. // Open connection
  288. act(() => {
  289. ws.open();
  290. });
  291. // Simulate print complete
  292. act(() => {
  293. ws.simulateMessage({
  294. type: 'print_complete',
  295. printer_id: 1,
  296. data: { status: 'completed' },
  297. });
  298. });
  299. // Advance timers to trigger debounced invalidation (3000ms delay + 500ms between each)
  300. await act(async () => {
  301. vi.advanceTimersByTime(4000);
  302. });
  303. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['archives'] });
  304. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['archiveStats'] });
  305. vi.useRealTimers();
  306. vi.unstubAllGlobals();
  307. });
  308. it('invalidates archives on archive_created message', async () => {
  309. vi.useFakeTimers();
  310. const { useWebSocket } = await import('../../hooks/useWebSocket');
  311. const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries');
  312. renderHook(() => useWebSocket(), {
  313. wrapper: createWrapper(queryClient),
  314. });
  315. const ws = await waitForWs();
  316. // Open connection
  317. act(() => {
  318. ws.open();
  319. });
  320. // Simulate archive created
  321. act(() => {
  322. ws.simulateMessage({
  323. type: 'archive_created',
  324. data: { id: 1, filename: 'test.3mf' },
  325. });
  326. });
  327. // Advance timers to trigger debounced invalidation (3000ms delay + 500ms between each)
  328. await act(async () => {
  329. vi.advanceTimersByTime(4000);
  330. });
  331. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['archives'] });
  332. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['archiveStats'] });
  333. vi.useRealTimers();
  334. vi.unstubAllGlobals();
  335. });
  336. it('invalidates archives on archive_updated message', async () => {
  337. vi.useFakeTimers();
  338. const { useWebSocket } = await import('../../hooks/useWebSocket');
  339. const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries');
  340. renderHook(() => useWebSocket(), {
  341. wrapper: createWrapper(queryClient),
  342. });
  343. const ws = await waitForWs();
  344. // Open connection
  345. act(() => {
  346. ws.open();
  347. });
  348. // Simulate archive updated (e.g., timelapse attached)
  349. act(() => {
  350. ws.simulateMessage({
  351. type: 'archive_updated',
  352. data: { id: 1, timelapse_attached: true },
  353. });
  354. });
  355. // Advance timers to trigger debounced invalidation (3000ms delay)
  356. await act(async () => {
  357. vi.advanceTimersByTime(4000);
  358. });
  359. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['archives'] });
  360. vi.useRealTimers();
  361. vi.unstubAllGlobals();
  362. });
  363. it('invalidates inventory queries on inventory_changed message', async () => {
  364. vi.useFakeTimers();
  365. const { useWebSocket } = await import('../../hooks/useWebSocket');
  366. const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries');
  367. renderHook(() => useWebSocket(), {
  368. wrapper: createWrapper(queryClient),
  369. });
  370. const ws = await waitForWs();
  371. act(() => {
  372. ws.open();
  373. });
  374. act(() => {
  375. ws.simulateMessage({ type: 'inventory_changed' });
  376. });
  377. await act(async () => {
  378. vi.advanceTimersByTime(5000);
  379. });
  380. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['inventory-spools'] });
  381. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['spoolman-inventory-spools'] });
  382. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['inventory-locations'] });
  383. vi.useRealTimers();
  384. vi.unstubAllGlobals();
  385. });
  386. it('handles missing_spool_assignment message without error', async () => {
  387. const { useWebSocket } = await import('../../hooks/useWebSocket');
  388. renderHook(() => useWebSocket(), {
  389. wrapper: createWrapper(queryClient),
  390. });
  391. const ws = await waitForWs();
  392. act(() => {
  393. ws.open();
  394. });
  395. // This test verifies that the hook properly handles missing_spool_assignment messages
  396. // without throwing an error. The actual toast display is tested via the UI.
  397. expect(() => {
  398. act(() => {
  399. ws.simulateMessage({
  400. type: 'missing_spool_assignment',
  401. printer_id: 7,
  402. printer_name: 'Printer B',
  403. missing_slots: [{ slot: 'A2' }, { slot: 'Ext-L' }],
  404. });
  405. });
  406. }).not.toThrow();
  407. vi.unstubAllGlobals();
  408. });
  409. it('shows an error toast when the billing kill switch stops a print', async () => {
  410. const { useWebSocket } = await import('../../hooks/useWebSocket');
  411. renderHook(() => useWebSocket(), {
  412. wrapper: createWrapper(queryClient),
  413. });
  414. const ws = await waitForWs();
  415. act(() => {
  416. ws.open();
  417. ws.simulateMessage({
  418. type: 'kill_switch_triggered',
  419. printer_id: 7,
  420. printer_name: 'Printer B',
  421. filename: 'foreign_job.3mf',
  422. });
  423. });
  424. const toast = screen.getByText('Billing kill switch stopped foreign_job.3mf on Printer B');
  425. expect(toast.parentElement).toHaveClass('bg-red-500/10');
  426. });
  427. it('handles spool_assignment_verified messages (success and failure) without error', async () => {
  428. const { useWebSocket } = await import('../../hooks/useWebSocket');
  429. renderHook(() => useWebSocket(), {
  430. wrapper: createWrapper(queryClient),
  431. });
  432. const ws = await waitForWs();
  433. act(() => {
  434. ws.open();
  435. });
  436. // #2582: verified (loaded), loaded-but-no-K-profile, and not-confirmed
  437. // all route to a toast — assert none of the branches throw.
  438. expect(() => {
  439. act(() => {
  440. ws.simulateMessage({
  441. type: 'spool_assignment_verified',
  442. printer_id: 3,
  443. printer_name: 'Printer A',
  444. slot: 'A1',
  445. verified: true,
  446. kprofile_applied: true,
  447. });
  448. ws.simulateMessage({
  449. type: 'spool_assignment_verified',
  450. printer_id: 3,
  451. printer_name: 'Printer A',
  452. slot: 'A1',
  453. verified: true,
  454. kprofile_applied: false,
  455. });
  456. ws.simulateMessage({
  457. type: 'spool_assignment_verified',
  458. printer_id: 3,
  459. printer_name: 'Printer A',
  460. slot: 'A1',
  461. verified: false,
  462. saw_tray: true,
  463. });
  464. });
  465. }).not.toThrow();
  466. vi.unstubAllGlobals();
  467. });
  468. it('ignores pong messages without error', async () => {
  469. const { useWebSocket } = await import('../../hooks/useWebSocket');
  470. const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries');
  471. renderHook(() => useWebSocket(), {
  472. wrapper: createWrapper(queryClient),
  473. });
  474. const ws = await waitForWs();
  475. // Open connection
  476. act(() => {
  477. ws.open();
  478. });
  479. // Simulate pong response
  480. act(() => {
  481. ws.simulateMessage({
  482. type: 'pong',
  483. });
  484. });
  485. // Should not invalidate any queries for pong
  486. expect(invalidateSpy).not.toHaveBeenCalled();
  487. });
  488. it('handles malformed JSON gracefully', async () => {
  489. const { useWebSocket } = await import('../../hooks/useWebSocket');
  490. renderHook(() => useWebSocket(), {
  491. wrapper: createWrapper(queryClient),
  492. });
  493. const ws = await waitForWs();
  494. // Open connection
  495. act(() => {
  496. ws.open();
  497. });
  498. // Simulate malformed message (should not throw)
  499. expect(() => {
  500. act(() => {
  501. if (ws.onmessage) {
  502. ws.onmessage(
  503. new MessageEvent('message', {
  504. data: 'not valid json{{{',
  505. })
  506. );
  507. }
  508. });
  509. }).not.toThrow();
  510. });
  511. it('handles unknown message types gracefully', async () => {
  512. const { useWebSocket } = await import('../../hooks/useWebSocket');
  513. const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries');
  514. renderHook(() => useWebSocket(), {
  515. wrapper: createWrapper(queryClient),
  516. });
  517. const ws = await waitForWs();
  518. // Open connection
  519. act(() => {
  520. ws.open();
  521. });
  522. // Simulate unknown message type
  523. expect(() => {
  524. act(() => {
  525. ws.simulateMessage({
  526. type: 'unknown_type',
  527. data: { foo: 'bar' },
  528. });
  529. });
  530. }).not.toThrow();
  531. expect(invalidateSpy).not.toHaveBeenCalled();
  532. });
  533. });
  534. /**
  535. * #2754 (reporter @mic4rd): live updates froze whenever the tab wasn't in
  536. * front, and caught up all at once on switching back.
  537. *
  538. * Two causes, fixed in two rounds. First the cache writes ran inside a
  539. * requestAnimationFrame, and a hidden tab gets no rendering opportunities —
  540. * the browser holds queued frame callbacks indefinitely rather than merely
  541. * throttling them. The rAF stub below is what makes those tests meaningful:
  542. * it hands back a handle and never invokes the callback, which is what a
  543. * real hidden tab does.
  544. *
  545. * Removing the frame callback did not close the report, because the 100ms
  546. * coalescing timer was still in the path and a hidden page's timers are
  547. * clamped to at best once a second — once a minute past five minutes hidden.
  548. * So the writes must not depend on a timer either while hidden, which is
  549. * what `writes without waiting on a timer` pins down. Note it deliberately
  550. * never advances the clock: a test that advances fake timers cannot tell a
  551. * throttled timer from a prompt one, which is exactly why the original tests
  552. * kept passing while the reporter's tab stayed frozen.
  553. */
  554. describe('hidden tab (#2754)', () => {
  555. let rafSpy: ReturnType<typeof vi.fn>;
  556. beforeEach(() => {
  557. // The shared test client sets gcTime: 0, which collects a query the
  558. // moment it has no observers — advancing timers past the 100ms
  559. // coalescing window would drop the entry we just wrote before we could
  560. // read it back. Nothing observes ['printerStatus', 1] here, so this
  561. // block needs a client that keeps unobserved data.
  562. queryClient = new QueryClient({
  563. defaultOptions: { queries: { retry: false, gcTime: Infinity } },
  564. });
  565. Object.defineProperty(document, 'hidden', { configurable: true, value: true });
  566. // Order matters: vi.useFakeTimers() fakes requestAnimationFrame as well
  567. // (backing it with the mock clock, so advanceTimersByTime would run it
  568. // and hide the very defect under test). Stub it afterwards so the
  569. // never-firing version is the one the hook sees.
  570. vi.useFakeTimers();
  571. rafSpy = vi.fn(() => 1);
  572. vi.stubGlobal('requestAnimationFrame', rafSpy);
  573. });
  574. afterEach(() => {
  575. vi.useRealTimers();
  576. Object.defineProperty(document, 'hidden', { configurable: true, value: false });
  577. });
  578. it('applies printer status to the query cache', async () => {
  579. const { useWebSocket } = await import('../../hooks/useWebSocket');
  580. renderHook(() => useWebSocket(), { wrapper: createWrapper(queryClient) });
  581. const ws = await waitForWs();
  582. act(() => ws.open());
  583. act(() => {
  584. ws.simulateMessage({
  585. type: 'printer_status',
  586. printer_id: 1,
  587. data: { state: 'RUNNING', progress: 42 },
  588. });
  589. });
  590. // Past the 100ms coalescing window.
  591. await act(async () => {
  592. vi.advanceTimersByTime(200);
  593. });
  594. // This is the key the tab-title/favicon progress reads
  595. // (usePrintProgressTitle) and nothing else.
  596. expect(queryClient.getQueryData(['printerStatus', 1])).toMatchObject({
  597. state: 'RUNNING',
  598. progress: 42,
  599. });
  600. expect(rafSpy).not.toHaveBeenCalled();
  601. });
  602. it('writes without waiting on a timer', async () => {
  603. const { useWebSocket } = await import('../../hooks/useWebSocket');
  604. renderHook(() => useWebSocket(), { wrapper: createWrapper(queryClient) });
  605. const ws = await waitForWs();
  606. act(() => ws.open());
  607. act(() => {
  608. ws.simulateMessage({
  609. type: 'printer_status',
  610. printer_id: 1,
  611. data: { state: 'RUNNING', progress: 42 },
  612. });
  613. });
  614. // No advanceTimersByTime: a hidden tab's timers are throttled to once a
  615. // second at best, so anything the title depends on has to have landed
  616. // already. Reintroduce the coalescing timer on this path and the cache
  617. // is still empty here.
  618. expect(queryClient.getQueryData(['printerStatus', 1])).toMatchObject({
  619. state: 'RUNNING',
  620. progress: 42,
  621. });
  622. });
  623. it('applies the newest value when several arrive before a frame would have run', async () => {
  624. const { useWebSocket } = await import('../../hooks/useWebSocket');
  625. renderHook(() => useWebSocket(), { wrapper: createWrapper(queryClient) });
  626. const ws = await waitForWs();
  627. act(() => ws.open());
  628. act(() => {
  629. ws.simulateMessage({ type: 'printer_status', printer_id: 1, data: { progress: 40 } });
  630. ws.simulateMessage({ type: 'printer_status', printer_id: 1, data: { progress: 41 } });
  631. });
  632. // Writing through per message must not resurrect an earlier one: the
  633. // pending map is drained on each flush, so a stale entry cannot be
  634. // re-applied over the newer value.
  635. expect(queryClient.getQueryData(['printerStatus', 1])).toMatchObject({ progress: 41 });
  636. });
  637. it('drains queued messages instead of wedging the queue', async () => {
  638. const { useWebSocket } = await import('../../hooks/useWebSocket');
  639. const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries');
  640. renderHook(() => useWebSocket(), { wrapper: createWrapper(queryClient) });
  641. const ws = await waitForWs();
  642. act(() => ws.open());
  643. // Everything other than printer_status goes through the message queue,
  644. // which used to stall with processingRef stuck true — messages then
  645. // piled up unbounded until the tab was shown again.
  646. act(() => {
  647. ws.simulateMessage({ type: 'print_complete', printer_id: 1, data: {} });
  648. });
  649. // 3s debounce, then the 500ms-apart stagger.
  650. await act(async () => {
  651. vi.advanceTimersByTime(4000);
  652. });
  653. expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ['archives'] });
  654. expect(rafSpy).not.toHaveBeenCalled();
  655. });
  656. });
  657. describe('visible tab still coalesces (#2754)', () => {
  658. /**
  659. * The counterpart to the hidden-tab block: the write-through is scoped to
  660. * a hidden tab on purpose. A visible one is painting, and the 100ms window
  661. * is what stops a burst of status messages turning into a render cascade —
  662. * so "just always write through" is not the simplification it looks like.
  663. */
  664. it('defers the write while the tab is visible', async () => {
  665. const { useWebSocket } = await import('../../hooks/useWebSocket');
  666. const client = new QueryClient({
  667. defaultOptions: { queries: { retry: false, gcTime: Infinity } },
  668. });
  669. vi.useFakeTimers();
  670. try {
  671. renderHook(() => useWebSocket(), { wrapper: createWrapper(client) });
  672. const ws = await waitForWs();
  673. act(() => ws.open());
  674. act(() => {
  675. ws.simulateMessage({
  676. type: 'printer_status',
  677. printer_id: 1,
  678. data: { state: 'RUNNING', progress: 42 },
  679. });
  680. });
  681. expect(client.getQueryData(['printerStatus', 1])).toBeUndefined();
  682. await act(async () => {
  683. vi.advanceTimersByTime(200);
  684. });
  685. expect(client.getQueryData(['printerStatus', 1])).toMatchObject({ progress: 42 });
  686. } finally {
  687. vi.useRealTimers();
  688. }
  689. });
  690. });
  691. describe('sendMessage', () => {
  692. it('sends JSON message when connected', async () => {
  693. const { useWebSocket } = await import('../../hooks/useWebSocket');
  694. const { result } = renderHook(() => useWebSocket(), {
  695. wrapper: createWrapper(queryClient),
  696. });
  697. const ws = await waitForWs();
  698. // Open connection
  699. act(() => {
  700. ws.open();
  701. });
  702. act(() => {
  703. result.current.sendMessage({ type: 'test', data: 'hello' });
  704. });
  705. expect(ws.send).toHaveBeenCalledWith(
  706. JSON.stringify({ type: 'test', data: 'hello' })
  707. );
  708. });
  709. it('does not send when disconnected', async () => {
  710. const { useWebSocket } = await import('../../hooks/useWebSocket');
  711. const { result } = renderHook(() => useWebSocket(), {
  712. wrapper: createWrapper(queryClient),
  713. });
  714. const ws = await waitForWs();
  715. // Don't open connection - still in CONNECTING state
  716. act(() => {
  717. result.current.sendMessage({ type: 'test' });
  718. });
  719. expect(ws.send).not.toHaveBeenCalled();
  720. });
  721. });
  722. describe('reconnection', () => {
  723. it('reconnects after connection closes', async () => {
  724. vi.useFakeTimers();
  725. const { useWebSocket } = await import('../../hooks/useWebSocket');
  726. renderHook(() => useWebSocket(), {
  727. wrapper: createWrapper(queryClient),
  728. });
  729. // GHSA-r2qv: connect() awaits a ws-token fetch before constructing
  730. // the WebSocket. Flush microtasks under fake timers so the await
  731. // resolves and MockWebSocket is pushed into wsInstances.
  732. await vi.advanceTimersByTimeAsync(0);
  733. const firstWs = wsInstances[wsInstances.length - 1]!;
  734. // Open connection
  735. act(() => {
  736. firstWs.open();
  737. });
  738. const instanceCountBefore = wsInstances.length;
  739. // Close connection
  740. act(() => {
  741. firstWs.close();
  742. });
  743. // Wait for reconnect timeout (3 seconds) + microtask flush for the
  744. // async connect() that the reconnect schedules.
  745. await vi.advanceTimersByTimeAsync(3000);
  746. // Should have created new WebSocket
  747. expect(wsInstances.length).toBe(instanceCountBefore + 1);
  748. expect(wsInstances[wsInstances.length - 1]).not.toBe(firstWs);
  749. vi.useRealTimers();
  750. });
  751. it('does NOT reconnect after an auth-rejection close (4401)', async () => {
  752. // Regression: a 4401 (ws-token invalid/expired or caller lacks
  753. // WEBSOCKET_CONNECT) used to reschedule connect() every 3s, spamming
  754. // /auth/ws-token forever. It must be terminal now.
  755. vi.useFakeTimers();
  756. const { useWebSocket } = await import('../../hooks/useWebSocket');
  757. renderHook(() => useWebSocket(), { wrapper: createWrapper(queryClient) });
  758. await vi.advanceTimersByTimeAsync(0);
  759. const firstWs = wsInstances[wsInstances.length - 1]!;
  760. act(() => {
  761. firstWs.open();
  762. });
  763. const instanceCountBefore = wsInstances.length;
  764. // Server rejects auth.
  765. act(() => {
  766. firstWs.simulateClose(4401);
  767. });
  768. // No reconnect even after the 3s window elapses.
  769. await vi.advanceTimersByTimeAsync(3000);
  770. expect(wsInstances.length).toBe(instanceCountBefore);
  771. vi.useRealTimers();
  772. });
  773. it('does NOT open a socket or reconnect when ws-token mint returns 403', async () => {
  774. // Mike/Forge's case: an authenticated user whose group lacks
  775. // WEBSOCKET_CONNECT. POST /auth/ws-token returns 403; the hook must NOT
  776. // fall through to a tokenless socket (server closes it 4401) and must NOT
  777. // enter the reconnect loop — it degrades to REST polling instead.
  778. vi.useFakeTimers();
  779. vi.stubGlobal(
  780. 'fetch',
  781. vi.fn(async () => ({
  782. ok: false,
  783. status: 403,
  784. statusText: 'Forbidden',
  785. headers: { get: () => null },
  786. json: async () => ({ detail: 'Insufficient permissions' }),
  787. })),
  788. );
  789. const { useWebSocket } = await import('../../hooks/useWebSocket');
  790. renderHook(() => useWebSocket(), { wrapper: createWrapper(queryClient) });
  791. // Flush the token-mint rejection, then let the (would-be) reconnect
  792. // window pass. No socket should ever be constructed.
  793. await vi.advanceTimersByTimeAsync(0);
  794. await vi.advanceTimersByTimeAsync(3000);
  795. expect(wsInstances.length).toBe(0);
  796. vi.useRealTimers();
  797. });
  798. it('does NOT reconnect when a close fires during unmount', async () => {
  799. // The provider unmounting (e.g. logout redirect) must not leave a
  800. // scheduled reconnect behind — the cleanup marks disposed before
  801. // close(), so the resulting onclose is a no-op.
  802. vi.useFakeTimers();
  803. const { useWebSocket } = await import('../../hooks/useWebSocket');
  804. const { unmount } = renderHook(() => useWebSocket(), {
  805. wrapper: createWrapper(queryClient),
  806. });
  807. await vi.advanceTimersByTimeAsync(0);
  808. const ws = wsInstances[wsInstances.length - 1]!;
  809. act(() => {
  810. ws.open();
  811. });
  812. const instanceCountBefore = wsInstances.length;
  813. // Unmount closes the socket, which fires onclose synchronously.
  814. act(() => {
  815. unmount();
  816. });
  817. await vi.advanceTimersByTimeAsync(3000);
  818. expect(wsInstances.length).toBe(instanceCountBefore);
  819. vi.useRealTimers();
  820. });
  821. it('cleans up on unmount', async () => {
  822. const { useWebSocket } = await import('../../hooks/useWebSocket');
  823. const { unmount } = renderHook(() => useWebSocket(), {
  824. wrapper: createWrapper(queryClient),
  825. });
  826. const ws = await waitForWs();
  827. // Open connection
  828. act(() => {
  829. ws.open();
  830. });
  831. unmount();
  832. expect(ws.close).toHaveBeenCalled();
  833. });
  834. });
  835. });