machine.test.ts

import { Machine } from "./machine";
import { Context, State } from "./models";

describe("Machine", () => {
  describe("Machine.create(state)", () => {
    test("it returns an instance of state machine", () => {
      const machine = Machine.create({ value: "saying.hello", context: {} });
      expect(machine).toBeInstanceOf(Machine);
    });
  });

  describe("machine.value", () => {
    test("provides the current state value of the machine", () => {
      const machine = Machine.create({ value: "saying.hello", context: {} });
      expect(machine.value).toEqual("saying.hello");
    });
  });

  describe("machine.context", () => {
    test("provides the current context of the machine", () => {
      const machine = Machine.create({ value: "saying.hello", context: { phrase: "howdy" } });
      expect(machine.context).toEqual({ phrase: "howdy" });
    });
  });

  describe("Sending messages", () => {
    describe("machine.send(message)", () => {
      test("allows a user to send messages used to update the machine", () => {
        const initialState = { value: "initial", context: {} };
        const machine = Machine.create(initialState);
        const state = machine.send({ event: "anything", payload: {} });
        expect(state).toEqual(initialState);
      });
    });

    describe("machine.registerUpdate(fn)", () => {
      let machine: Machine<
        "saying.hello" | "saying.goodbye",
        { phrase: string },
        { event: "poke"; payload: {} } | { event: "sayGoodbye"; payload: { phrase?: string } }
      >;

      beforeEach(() => {
        machine = Machine.create({
          value: "saying.hello" as const,
          context: { phrase: "hi there!" },
        });
      });

      afterEach(() => {
        machine = undefined as any;
      });

      test("accepts a function used interpret messages", () => {
        machine.registerUpdate(({ event, payload }, state) => {
          switch (event) {
            case "sayGoodbye":
              return {
                value: "saying.goodbye",
                context: { ...state.context, phrase: payload.phrase ?? "goodbye" },
              };
            default:
              return state;
          }
        });

        machine.send({ event: "sayGoodbye", payload: {} });

        expect(machine.value).toEqual("saying.goodbye");
        expect(machine.context).toEqual({ phrase: "goodbye" });
      });

      test("a bad update cannot take down the system", () => {
        let history: { value: string; context: { phrase: string } }[] = [];

        machine.registerUpdate(({ event, payload }, state) => {
          if (event === "sayGoodbye") {
            return { value: "saying.goodbye", context: { phrase: payload.phrase ?? "whatevs" } };
          } else {
            return state;
          }
        });

        machine.registerUpdate(({ event }, state) => {
          if (event === "poke") {
            throw Error("OOOOOOOOOOOO");
          } else {
            return state;
          }
        });

        machine.subscribe(async (_, state) => {
          history.push(state);
        });

        machine.send({ event: "sayGoodbye", payload: { phrase: "later" } });
        machine.send({ event: "poke", payload: {} });
        machine.send({ event: "sayGoodbye", payload: { phrase: "buh bye" } });

        expect(history).toEqual([
          { value: "saying.goodbye", context: { phrase: "later" } },
          { value: "saying.goodbye", context: { phrase: "later" } },
          { value: "saying.goodbye", context: { phrase: "buh bye" } },
        ]);
      });
    });
  });

  describe("Subscribing to updates", () => {
    let machine: Machine<
      "counting",
      { count: number },
      { event: "increment"; payload: {} } | { event: "decrement"; payload: {} }
    >;

    beforeEach(() => {
      machine = Machine.create({ value: "counting", context: { count: 0 } });

      machine.registerUpdate(({ event }, { context, ...state }) => ({
        ...state,
        context: { count: context.count + (event === "increment" ? 1 : -1) },
      }));
    });

    afterEach(() => {
      machine = undefined as any;
    });

    describe("machine.subscribe()", () => {
      test("a user can subscribe to state updates", () => {
        let history: { count: number }[] = [];

        machine.subscribe(async (_, state) => {
          history.push(state.context);
        });

        machine.send({ event: "increment", payload: {} });
        machine.send({ event: "increment", payload: {} });
        machine.send({ event: "decrement", payload: {} });
        machine.send({ event: "increment", payload: {} });
        machine.send({ event: "decrement", payload: {} });

        expect(history).toEqual([
          { count: 1 },
          { count: 2 },
          { count: 1 },
          { count: 2 },
          { count: 1 },
        ]);
      });

      test("rogue subscribers cannot bring down the machine", () => {
        let history: Context[] = [];

        const good = jest.fn(async (_: unknown, state: State<string, Context>) => {
          history.push(state.context);
        });
        const bad = jest.fn(async () => {
          throw Error("Watch out!");
        });

        machine.subscribe(good);
        machine.subscribe(bad);

        machine.send({ event: "increment", payload: {} });
        machine.send({ event: "increment", payload: {} });
        machine.send({ event: "increment", payload: {} });

        expect(machine.context).toEqual({ count: 3 });
        expect(history).toEqual([{ count: 1 }, { count: 2 }, { count: 3 }]);
        expect(good).toHaveBeenCalledTimes(3);
        expect(bad).toHaveBeenCalledTimes(3);
      });
    });
  });

  describe("handling side effects", () => {
    let machine: Machine<
      "idle" | "loading" | "success" | "failure",
      { data?: string; error?: string },
      | { event: "request"; payload: { what: string } }
      | { event: "succeed"; payload: { data: string } }
      | { event: "fail"; payload: { error: string } }
    >;

    let history: { value: string; context: {} }[] = [];

    beforeEach(() => {
      machine = Machine.create({
        value: "idle",
        context: {},
      });

      machine.registerUpdate(({ event, payload }, state) => {
        switch (event) {
          case "request":
            return { value: "loading", context: {} };
          case "succeed":
            return { value: "success", context: { ...state.context, data: payload.data } };
          case "fail":
            return { value: "failure", context: { ...state.context, error: payload.error } };
          default:
            return state;
        }
      });

      machine.subscribe(async (_, state) => {
        history.push(state);
      });
    });

    afterEach(() => {
      machine = undefined as any;
      history = [];
    });

    describe("machine.registerEffect(effect)", () => {
      test("effects can update state in response to a message", () => {
        machine.registerEffect(async ({ event, payload }, send) => {
          if (event === "request" && payload.what !== "oops") {
            const data = `${payload.what.toUpperCase()}!!`;
            send({ event: "succeed", payload: { data } });
          } else if (event === "request") {
            const error = "UH OH";
            send({ event: "fail", payload: { error } });
          }
        });

        machine.send({ event: "request", payload: { what: "wow" } });
        machine.send({ event: "request", payload: { what: "oops" } });

        expect(history).toEqual([
          { value: "loading", context: {} },
          { value: "success", context: { data: "WOW!!" } },
          { value: "loading", context: {} },
          { value: "failure", context: { error: "UH OH" } },
        ]);
      });

      test("effects can not bring down the system", () => {
        machine.registerEffect(async ({ event, payload }, send) => {
          if (event === "request" && payload.what === "oops") {
            throw Error("oh no");
          } else if (event === "request") {
            const data = `${payload.what.toUpperCase()}!!`;
            send({ event: "succeed", payload: { data } });
          }
        });

        machine.send({ event: "request", payload: { what: "wow" } });
        machine.send({ event: "request", payload: { what: "oops" } });
        machine.send({ event: "request", payload: { what: "cool" } });

        expect(history).toEqual([
          { value: "loading", context: {} },
          { value: "success", context: { data: "WOW!!" } },
          { value: "loading", context: {} },
          { value: "loading", context: {} },
          { value: "success", context: { data: "COOL!!" } },
        ]);
      });
    });
  });
});

machine.ts

import type {
  Callback,
  Context as TContext,
  Effect,
  Message as TMessage,
  State as TState,
  Update,
} from "./models";

const reportError = console.error.bind(console); // TODO: Extract and report to datadog

export class Machine<
  Value extends string,
  Context extends TContext,
  Message extends TMessage,
  State extends TState<Value, Context> = TState<Value, Context>,
> {
  private updates: Map<Symbol, Update<Message, State>> = new Map();
  private subscribers: Map<Symbol, Callback<Message, State>> = new Map();
  private effects: Map<Symbol, Effect<Message, State>> = new Map();

  private constructor(private state: State) {}

  get value() {
    return this.state.value;
  }

  get context() {
    return this.state.context;
  }

  static create = <
    V extends string,
    C extends TContext,
    M extends TMessage,
    S extends TState<V, C> = TState<V, C>,
  >(
    state: S,
  ) => new Machine<V, C, M>(state);

  private broadcast = async (message: Message, state: State) => {
    const subscribers = [...this.subscribers.values()].map((callback) => callback(message, state));
    Promise.allSettled(subscribers); // TODO: Error reporting for failed callbacks
  };

  private runEffects = async (message: Message) => {
    const effects = [...this.effects.values()].map((effect) => effect(message, this.send));
    Promise.allSettled(effects); // TODO: Error reporting for failed effects
  };

  private update = (message: Message, state: State): State => {
    const updates = [...this.updates.values()];

    return updates.reduce((stateAcc, fn) => {
      let updatedState = stateAcc;
      try {
        updatedState = fn(message, stateAcc);
      } catch (err) {
        reportError(err);
      } finally {
        return updatedState;
      }
    }, state);
  };

  registerEffect = (effect: Effect<Message, State>) => {
    this.effects.set(Symbol(), effect);
    return this;
  };

  registerUpdate = (update: Update<Message, State>) => {
    this.updates.set(Symbol(), update);
    return this;
  };

  subscribe = (callback: Callback<Message, State>) => {
    this.subscribers.set(Symbol(), callback);
    return this;
  };

  send = (message: Message) => {
    this.state = this.update(message, this.state);
    this.broadcast(message, this.state);
    this.runEffects(message);
    return this.state;
  };
}

model.ts

export type Context = Record<string, unknown>;

export type State<V extends string, C extends Context> = {
  value: V;
  context: C;
};

export type Message = {
  event: string;
  payload: Record<string, unknown>;
};

export type Update<M, S extends State<string, Context>> = (message: M, state: S) => S;

export type Effect<M, S extends State<string, Context>> = (
  message: M,
  send: (message: M) => S,
) => Promise<void>;

export type Callback<M extends Message, S extends State<string, Context>> = (
  message: M,
  state: S,
) => Promise<void>;