How to test xstate machines invoking async code in Jest

I’m writing a complex state machine which invokes a number of promises. Those promises are calls to an external API or to the database.

The tests use msw to mock API calls, and a test database which is seeded and torn down before and after each test.

Since my state machine moves from one state to another sequentially, I opted to test that my implementation was correct at each step. Here’s an example, adapted from the xstate docs:

describe('call_locations_api', () => {
  it('should save locations to context on success and go to compare step', (done) => {
    interpret(updateLocationsMachine)
      .onTransition((state) => {
        if (state.matches('compare')) {
          expect(state.context.locations).toStrictEqual(locations);
          done();
        }
      })
      .start();
  });
});

I actually used an async/await variation, but for the sake of this article I’m using the done() callback as presented in the docs. The result is the same.

The problem

While tests fail or succeed as expected, the machine continues until it reaches a final state, logging a bunch of database errors to the console or worst, making API calls after the msw listener closes.

Calling done() only ends the test, not the machine itself.

The solution

Hidden under the “Interpreting machines” page of the xstate docs is a section about waitFor, which contains the following snippet:

const myFunc = async () => {
  const actor = interpret(machine).start();

  const doneState = await waitFor(actor, (state) => state.matches('done'));

  console.log(doneState.value); // 'done'
};

Huh, this is quite similar to the test code. Two key differences:

1. We keep a reference to the interpreted machine (actor).
2. The asynchronous nature of the code is explicit.

Adapting that to a test function, we can refactor our code:

import { waitFor } from 'xstate/lib/waitFor';

describe('call_locations_api', () => {
  it('should save locations to context on success and go to compare step', async () => {
    const actor = interpret(updateLocationsMachine).start();

    const state = await waitFor(actor, (state) => state.matches('compare'));

    expect(state.context.locations).toStrictEqual(locations);

    actor.stop();
  });
});

Keeping a reference to the machine (actor) allows us to stop it at the end of the test.

This solves the problem by preventing the state machine from persisting and, as a bonus, it makes the code more readable by extracting the assertion from the interpret block.