machineRegistry.ts

import { StateMachine, State } from "xstate";

type RegisteredMachine = {
  machine: StateMachine<any, any, any>;
  state?: State<any, any, any>;
};
const machineRegistrations: Record<string, RegisteredMachine> = {};

export function registerMachine(
  machine: StateMachine<any, any, any>,
  state?: State<any, any, any>
): void {
  machineRegistrations[machine.config.id!] = {
    machine,
    state,
  };
}

export function hasMachineRegistered(id: string): boolean {
  return id in machineRegistrations;
}

export function getMachine(id: string): RegisteredMachine {
  return machineRegistrations[id];
}

export function removeMachine(id: string): void {
  delete machineRegistrations[id];
}

useMachine.ts

import { useMachine as useMachineXState } from "@xstate/react";
import { ComponentType } from "react";
import { EventObject, Typestate, StateSchema, StateMachine, Interpreter } from "xstate";

import { getMachine, hasMachineRegistered } from "../../storybook/machineRegistry";
import { makeMachineServiceMatcher, MatchProps } from "../makeMachineMatcher";

export function useMachine<
  TContext,
  TStateSchema extends StateSchema,
  TEvent extends EventObject,
  TTypestate extends Typestate<TContext>,
  TInterpreter extends Interpreter<TContext, TStateSchema, TEvent, TTypestate> = Interpreter<
    TContext,
    TStateSchema,
    TEvent,
    TTypestate
  >
>(
  machine: StateMachine<TContext, TStateSchema, TEvent, TTypestate>
): [
  ComponentType<MatchProps<TContext, TStateSchema, TEvent, TTypestate, TTypestate["value"]>>,
  TInterpreter["state"],
  TInterpreter["send"],
  TInterpreter
] {
  let machineToStart = machine;
  let initialState = undefined;
  // All our machines have an ID set
  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  const machineId = machine.config.id!;

  if (hasMachineRegistered(machineId)) {
    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
    const { machine: machineOverride, state } = getMachine(machineId);

    machineToStart = machineOverride;
    initialState = state;
  }

  const [current, send, service] = useMachineXState(machineToStart, {
    immediate: true,
    state: initialState,
  });
  const Match = makeMachineServiceMatcher(service);

  return [
    Match,
    (current as unknown) as TInterpreter["state"],
    send,
    (service as unknown) as TInterpreter,
  ];
}

withMachine.tsx

import { makeDecorator } from "@storybook/addons";
import React, { useEffect } from "react";

import { $YesReallyAny } from "../";

import { registerMachine, removeMachine } from "./machineRegistry";

function withMachineContext(storyFn: () => $YesReallyAny, machineId: string): React.ComponentType {
  return () => {
    useEffect(() => {
      return () => {
        removeMachine(machineId);
      };
    }, []);

    return <>{storyFn()}</>;
  };
}

export const withMachine = makeDecorator({
  name: "withMachine",
  parameterName: "withMachine",
  wrapper: (storyFn, context, { options }) => {
    const { machine, state } = options;
    registerMachine(machine, state);

    return React.createElement(withMachineContext(() => storyFn(context), machine.config.id));
  },
});