12.7 KB430 lines
Blame
1/**
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8import type {Atom} from 'jotai';
9import type {AtomFamilyWeak} from '../jotaiUtils';
10import type {Platform} from '../platform';
11import type {LocalStorageName} from '../types';
12
13import {act, render} from '@testing-library/react';
14import {List} from 'immutable';
15import {Provider, atom, createStore, useAtomValue} from 'jotai';
16import {StrictMode} from 'react';
17import {SelfUpdate} from 'shared/immutableExt';
18import {gc, nextTick} from 'shared/testUtils';
19import {
20 atomFamilyWeak,
21 atomResetOnDepChange,
22 lazyAtom,
23 localStorageBackedAtomFamily,
24 readAtom,
25 useAtomGet,
26 useAtomHas,
27 writeAtom,
28} from '../jotaiUtils';
29
30class Foo extends SelfUpdate<List<number>> {}
31
32const testFooAtom = atom<Foo>(new Foo(List()));
33
34describe('Jotai compatibility', () => {
35 it('does not freeze SelfUpdate types', () => {
36 function MyTestComponent() {
37 const foo2 = useAtomValue(testFooAtom);
38 expect(Object.isSealed(foo2)).toBe(false);
39 expect(Object.isFrozen(foo2)).toBe(false);
40
41 return null;
42 }
43 render(
44 <Provider>
45 <MyTestComponent />
46 </Provider>,
47 );
48 });
49});
50
51describe('lazyAtom', () => {
52 it('returns sync load() value', () => {
53 const a = lazyAtom(() => 1, 2);
54 expect(readAtom(a)).toBe(1);
55 });
56
57 it('returns fallback value and sets async load() value', async () => {
58 const a = lazyAtom(() => Promise.resolve(1), 2);
59 expect(readAtom(a)).toBe(2);
60 await nextTick();
61 expect(readAtom(a)).toBe(1);
62 });
63
64 it('can depend on another atom', () => {
65 const a = atom<number | undefined>(undefined);
66 const b = lazyAtom(() => 2, 2, a);
67
68 expect(readAtom(a)).toBe(undefined);
69
70 // Reading `b` triggers updating `a`.
71 expect(readAtom(b)).toBe(2);
72 expect(readAtom(a)).toBe(2);
73
74 // If `a` is updated to be not `undefined`, `b` will be the same value.
75 writeAtom(a, 3);
76 expect(readAtom(b)).toBe(3);
77
78 // Updating `b` updates `a` too.
79 writeAtom(b, 4);
80 expect(readAtom(a)).toBe(4);
81
82 // If `a` is updated to be `undefined`, `b` will be the fallback value.
83 writeAtom(a, undefined);
84 expect(readAtom(b)).toBe(2);
85 });
86});
87
88describe('atomFamilyWeak', () => {
89 let recalcFamilyCount = 0;
90 let recalcAtomCount = 0;
91 const family = atomFamilyWeak((key: number) => {
92 recalcFamilyCount += 1;
93 return atom(_get => {
94 recalcAtomCount += 1;
95 return key;
96 });
97 });
98
99 let store = createStore();
100
101 function TestComponent({k, f}: {k: number; f?: AtomFamilyWeak<number, Atom<number>>}) {
102 const value = useAtomValue((f ?? family)(k));
103 return <div>{value}</div>;
104 }
105 function TestApp({
106 keys,
107 family,
108 }: {
109 keys: Array<number>;
110 family?: AtomFamilyWeak<number, Atom<number>>;
111 }) {
112 return (
113 <Provider store={store}>
114 {keys.map(k => (
115 <TestComponent k={k} key={k} f={family} />
116 ))}
117 </Provider>
118 );
119 }
120
121 beforeEach(() => {
122 store = createStore();
123 family.clear();
124 recalcFamilyCount = 0;
125 recalcAtomCount = 0;
126 });
127
128 it('with "strong" cache enabled (default), unmount/remount skips re-calculate', async () => {
129 const rendered = render(<TestApp keys={[1, 2, 3]} />);
130 expect(recalcFamilyCount).toBe(3);
131 expect(recalcAtomCount).toBe(3);
132
133 rendered.rerender(<TestApp keys={[1, 2, 3]} />);
134 expect(recalcFamilyCount).toBe(3);
135 expect(recalcAtomCount).toBe(3);
136
137 // After unmount, part of the atomFamily cache can still be used when re-mounting.
138 rendered.unmount();
139 await gc();
140
141 render(<TestApp keys={[3, 2, 1]} />);
142 expect(recalcFamilyCount).toBeLessThan(6);
143 expect(recalcAtomCount).toBeLessThan(6);
144 });
145
146 it('with "strong" cache disabled, unmount/remount might re-calculate', async () => {
147 let count = 0;
148 const family = atomFamilyWeak(
149 (key: number) => {
150 count += 1;
151 return atom(_get => key);
152 },
153 {useStrongCache: false},
154 );
155
156 const rendered = render(<TestApp keys={[1, 2, 3]} family={family} />);
157 expect(count).toBe(3);
158
159 rendered.unmount();
160 await gc();
161
162 render(<TestApp keys={[1, 2, 3]} family={family} />);
163 expect(count).toBe(6);
164 });
165
166 it('"cleanup" does not clean in-use states', () => {
167 const rendered = render(<TestApp keys={[1, 2, 3]} />);
168 expect(recalcFamilyCount).toBe(3);
169 expect(recalcAtomCount).toBe(3);
170 family.cleanup();
171
172 // re-render can still use cached state after "cleanup" (count remains the same).
173 rendered.rerender(<TestApp keys={[1, 2, 3]} />);
174 expect(recalcFamilyCount).toBe(3);
175 expect(recalcAtomCount).toBe(3);
176 });
177
178 it('"cleanup" can release memory', async () => {
179 const rendered = render(<TestApp keys={[1, 2, 3]} />);
180 expect(recalcFamilyCount).toBe(3);
181 expect(recalcAtomCount).toBe(3);
182 rendered.unmount();
183
184 await gc();
185 family.cleanup();
186 await gc();
187 family.cleanup();
188
189 // umount, then re-render will recalculate all atoms (count increases).
190 render(<TestApp keys={[3, 2, 1]} />);
191 expect(recalcFamilyCount).toBe(6);
192 expect(recalcAtomCount).toBe(6);
193 });
194
195 it('"clear" releases memory', () => {
196 const rendered = render(<TestApp keys={[1, 2, 3]} />);
197 expect(recalcFamilyCount).toBe(3);
198 expect(recalcAtomCount).toBe(3);
199 family.clear();
200
201 // re-render will recalculate all atoms.
202 rendered.rerender(<TestApp keys={[1, 2, 3]} />);
203 expect(recalcFamilyCount).toBe(6);
204 expect(recalcAtomCount).toBe(6);
205 rendered.unmount();
206 family.clear();
207
208 // umount, then re-render will recalculate all atoms (count increases).
209 render(<TestApp keys={[3, 2, 1]} />);
210 expect(recalcFamilyCount).toBe(9);
211 expect(recalcAtomCount).toBe(9);
212 });
213
214 it('"cleanup" runs automatically to reduce cache size', async () => {
215 const N = 10;
216 const M = 30;
217
218 // Render N items.
219 const rendered = render(<TestApp keys={Array.from({length: N}, (_, i) => i)} />);
220
221 // Umount to drop references to the atoms.
222 rendered.unmount();
223
224 // Force GC to run to stabilize the test.
225 await gc();
226
227 // After GC, render M items with different keys.
228 // This would trigger `family.cleanup` transparently.
229 render(<TestApp keys={Array.from({length: M}, (_, i) => N + i)} />);
230
231 // Neither of the caches should have N + M items (which means no cleanup).
232 expect(family.weakCache.size).toBeLessThan(N + M);
233 expect(family.strongCache.size).toBeLessThan(N + M);
234 });
235
236 it('provides debugLabel', () => {
237 const family = atomFamilyWeak((v: string) => atom(v));
238 family.debugLabel = 'prefix1';
239 const atom1 = family('a');
240 expect(atom1.debugLabel).toBe('prefix1:a');
241 });
242});
243
244describe('useAtomGet and useAtomSet', () => {
245 const initialMap = new Map([
246 ['a', 1],
247 ['b', 2],
248 ['c', 3],
249 ]);
250 const initialSet = new Set(['a', 'b']);
251
252 // Render an App, change the map and set, check what
253 // Runs a test and report re-render and atom states.
254 // insertMap specifies changes to the map (initially {a: 1, b: 2, c: 3}).
255 // changeSet specifies changes to the set (initially {a, b}).
256 function findRerender(props: {
257 insertMap?: Iterable<[string, number]>;
258 replaceSet?: Iterable<string>;
259 }): Array<string> {
260 // container types
261 const map = atom<Map<string, number>>(initialMap);
262 const set = atom<Set<string>>(initialSet);
263 const rerenderKeys = new Set<string>();
264
265 // test UI components
266 function Item({k}: {k: string}) {
267 const mapValue = useAtomGet(map, k);
268 const setValue = useAtomHas(set, k);
269 rerenderKeys.add(k);
270 return (
271 <span>
272 {mapValue} {setValue}
273 </span>
274 );
275 }
276
277 const store = createStore();
278
279 function TestApp({keys}: {keys: Array<string>}) {
280 return (
281 <StrictMode>
282 <Provider store={store}>
283 {keys.map(k => (
284 <Item k={k} key={k} />
285 ))}
286 </Provider>
287 </StrictMode>
288 );
289 }
290
291 const keys = ['a', 'b', 'c'];
292 render(<TestApp keys={keys} />);
293
294 rerenderKeys.clear();
295
296 const {insertMap, replaceSet} = props;
297
298 act(() => {
299 if (insertMap) {
300 store.set(map, oldMap => new Map([...oldMap, ...insertMap]));
301 }
302 if (replaceSet) {
303 const newSet = new Set([...replaceSet]);
304 store.set(set, newSet);
305 }
306 });
307
308 return [...rerenderKeys];
309 }
310
311 it('avoids re-rendering with changing to unrelated keys', () => {
312 expect(findRerender({insertMap: [['unrelated-key', 3]]})).toEqual([]);
313 expect(findRerender({replaceSet: [...initialSet, 'unrelated-key']})).toEqual([]);
314 });
315
316 it('only re-render changed items', () => {
317 const replaceSet = [...initialSet, 'c']; // add 'c' to the set.
318 expect(findRerender({insertMap: [['b', 5]]})).toEqual(['b']);
319 expect(findRerender({replaceSet})).toEqual(['c']);
320 expect(findRerender({insertMap: [['b', 5]], replaceSet})).toEqual(['b', 'c']);
321 });
322});
323
324describe('atomResetOnDepChange', () => {
325 it('works like a primitive atom', () => {
326 const depAtom = atom(0);
327 const testAtom = atomResetOnDepChange(1, depAtom);
328 const doubleAtom = atom(get => get(testAtom) * 2);
329 expect(readAtom(doubleAtom)).toBe(2);
330 expect(readAtom(testAtom)).toBe(1);
331 writeAtom(testAtom, 2);
332 expect(readAtom(doubleAtom)).toBe(4);
333 expect(readAtom(testAtom)).toBe(2);
334 });
335
336 it('gets reset on dependency change', () => {
337 const depAtom = atom(0);
338 const testAtom = atomResetOnDepChange(1, depAtom);
339
340 writeAtom(testAtom, 2);
341
342 // Change depAtom should reset testAtom.
343 writeAtom(depAtom, 10);
344 expect(readAtom(testAtom)).toBe(1);
345
346 // Set depAtom to the same value does not reset testAtom.
347 writeAtom(testAtom, 3);
348 writeAtom(depAtom, 10);
349 expect(readAtom(testAtom)).toBe(3);
350 });
351});
352
353describe('localStorageBackedAtomFamily', () => {
354 function setupTestPlatform<T>(initialValues: Record<string, T>): Platform {
355 const state = {...initialValues};
356 const mockPlatform = {
357 getAllPersistedState: jest.fn().mockImplementation((): Record<string, T> => {
358 return state;
359 }),
360 getPersistedState: jest.fn().mockImplementation((key: string): T => {
361 return state[key];
362 }),
363 setPersistedState: jest.fn().mockImplementation((key: string, value: T) => {
364 state[key] = value;
365 }),
366 } as Partial<Platform> as Platform;
367 return mockPlatform;
368 }
369
370 const testKey = 'test_' as LocalStorageName;
371
372 it('loads initial state from storage', () => {
373 const mockPlatform = setupTestPlatform({test_a: {data: 1, date: Date.now()}});
374 const family = localStorageBackedAtomFamily(testKey, () => 0, 1, mockPlatform);
375
376 expect(readAtom(family('a'))).toEqual(1);
377 expect(readAtom(family('b'))).toEqual(0);
378 });
379
380 it('evicts old initial state from storage', () => {
381 const mockPlatform = setupTestPlatform({
382 test_a: {data: 1, date: 0},
383 test_b: {data: 2, date: Date.now()},
384 });
385 const family = localStorageBackedAtomFamily(testKey, () => 0, 1, mockPlatform);
386
387 expect(readAtom(family('a'))).toEqual(0);
388 expect(readAtom(family('b'))).toEqual(2);
389 });
390
391 it('writes to persisted state', () => {
392 const mockPlatform = setupTestPlatform<number>({});
393 const family = localStorageBackedAtomFamily(testKey, (): number => 0, 1, mockPlatform);
394
395 expect(readAtom(family('a'))).toEqual(0);
396 writeAtom(family('a'), 2);
397 expect(readAtom(family('a'))).toEqual(2);
398 expect(mockPlatform.setPersistedState).toHaveBeenCalledWith('test_a', {
399 data: 2,
400 date: expect.any(Number),
401 });
402 });
403
404 it("gets the latest data even after writing and then being gc'd", () => {
405 const mockPlatform = setupTestPlatform({test_a: {data: 1, date: Date.now()}});
406 const family = localStorageBackedAtomFamily(testKey, (): number => 0, 1, mockPlatform);
407
408 expect(readAtom(family('a'))).toEqual(1);
409 writeAtom(family('a'), 2);
410 expect(readAtom(family('a'))).toEqual(2);
411 family.clear();
412 expect(readAtom(family('a'))).toEqual(2);
413 });
414
415 it('writing undefined clears from persisted storage', () => {
416 const mockPlatform = setupTestPlatform({test_a: {data: 1, date: Date.now()}});
417 const family = localStorageBackedAtomFamily(testKey, () => 0, 1, mockPlatform);
418
419 // type check that although we can SET to undefined, we can't READ undefined
420 const returnTypeShouldBeNonNull: number = readAtom(family('a'));
421 expect(returnTypeShouldBeNonNull).toEqual(1);
422
423 writeAtom(family('a'), undefined);
424
425 expect(mockPlatform.setPersistedState).toHaveBeenCalledWith('test_a', undefined);
426 expect(mockPlatform.getPersistedState('test_a')).toBeUndefined(); // nothing is stored
427 expect(readAtom(family('a'))).toEqual(0); // read still gives default value
428 });
429});
430