How to Check for Memory Leaks in Unit Tests

Hi,

In my previous post, I briefly mentioned that Jotai has unit tests to check WeakMap behavior.

Jotai Memory Leak Tests

Let’s take a look at the test code:

import LeakDetector from 'jest-leak-detector';
import { expect, test } from 'vitest';

test('memory leaks with one atom', async () => {
  const store = createStore();
  let objAtom = atom({});
  const detector = new LeakDetector(store.get(objAtom));
  objAtom = undefined;
  await Promise.resolve();
  expect(await detector.isLeaking()).toBe(false);
});

Looks pretty simple, right? Now, let’s try the same concept without Jotai.

test('memory leaks with map', async () => {
  const store = new Map();
  let obj = {};
  store.set(obj, {});
  const detector = new LeakDetector(store.get(obj));
  obj = undefined;
  await Promise.resolve();
  expect(await detector.isLeaking()).toBe(false);
});
  
test('memory leaks with weakmap', async () => {
  const store = new WeakMap();
  let obj = {};
  store.set(obj, {});
  const detector = new LeakDetector(store.get(obj));
  obj = undefined;
  await Promise.resolve();
  expect(await detector.isLeaking()).toBe(false);
});

When you run these tests, the first one fails while the second one passes:

   × memory leaks with map 10ms
     → expected true to be false // Object.is equality
   ✓ memory leaks with weakmap

How jest-leak-detector Works

Finally, let’s take a look at how jest-leak-detector is implemented. It is part of the Jest repository but does not depend on Jest itself.

Source Code

The implementation is just 77 lines of code. Essentially, it uses FinalizationRegistry to track object finalization:

this._finalizationRegistry = new FinalizationRegistry(() => {
  this._isReferenceBeingHeld = false;
});
this._finalizationRegistry.register(value, undefined);
this._isReferenceBeingHeld = true;

Additionally, it includes a small hack to expose gc:

setFlagsFromString('--expose-gc');
runInNewContext('gc')();

While implementing this from scratch is not too difficult, I use jest-leak-detector because it works well with Vitest.

Hope you find this useful.

Happy coding.