From 260520893207af314a90f1f3483b69b4b70991e9 Mon Sep 17 00:00:00 2001
From: Thomas Reynolds <me@tdreyno.com>
Date: Mon, 18 Apr 2022 16:15:47 -0700
Subject: [PATCH] fix: onContextChange for updates

---
 src/__tests__/boundActions.spec.ts    |  2 +-
 src/__tests__/onContextChange.spec.ts | 30 ++++++++++++++++++---------
 src/runtime.ts                        |  7 +------
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/src/__tests__/boundActions.spec.ts b/src/__tests__/boundActions.spec.ts
index c6acd573..21c56830 100644
--- a/src/__tests__/boundActions.spec.ts
+++ b/src/__tests__/boundActions.spec.ts
@@ -49,6 +49,6 @@ describe("Bound actions", () => {
     ])
 
     expect(runtime.currentState().data).toBe(36)
-    expect(onChange).toHaveBeenCalledTimes(5)
+    expect(onChange).toHaveBeenCalledTimes(10)
   })
 })
diff --git a/src/__tests__/onContextChange.spec.ts b/src/__tests__/onContextChange.spec.ts
index ce829ee3..71ec8c84 100644
--- a/src/__tests__/onContextChange.spec.ts
+++ b/src/__tests__/onContextChange.spec.ts
@@ -1,6 +1,7 @@
+/* eslint-disable @typescript-eslint/no-unsafe-assignment */
 import { ActionCreatorType, Enter, createAction, enter } from "../action"
+import { Context, createInitialContext } from "../context"
 
-import { createInitialContext } from "../context"
 import { createRuntime } from "../runtime"
 import { noop } from "../effect"
 import { state } from "../state"
@@ -27,28 +28,37 @@ describe("onContextChange", () => {
     expect(onChange).toHaveBeenCalledTimes(1)
   })
 
-  test("should run callback once on update", async () => {
+  test("should run callback on update", async () => {
     const trigger = createAction("Trigger")
     type Trigger = ActionCreatorType<typeof trigger>
 
-    const A = state<Enter | Trigger, string>(
+    const A = state<Enter | Trigger, number>(
       {
-        Enter: noop,
-        Trigger: (name, _, { update }) => update(name + name),
+        Enter: (n, _, { update }) => [update(n + 1), trigger()],
+        Trigger: async (n, _, { update }) => update(n + 1),
       },
       { name: "A" },
     )
 
-    const context = createInitialContext([A("Test")])
+    const context = createInitialContext([A(1)])
 
     const runtime = createRuntime(context, ["Trigger"])
 
-    const onChange = jest.fn()
+    let i = 0
+    const onChange = jest.fn((context: Context) => {
+      const { data } = context.currentState
 
-    runtime.onContextChange(onChange)
+      if (i++ == 0) {
+        expect(data).toBe(2)
+      } else {
+        expect(data).toBe(3)
+      }
+    })
 
-    await runtime.run(trigger())
+    expect.assertions(3)
 
-    expect(onChange).toHaveBeenCalledTimes(1)
+    runtime.onContextChange(onChange)
+
+    await runtime.run(enter())
   })
 })
diff --git a/src/runtime.ts b/src/runtime.ts
index d328e47c..2c6ab404 100644
--- a/src/runtime.ts
+++ b/src/runtime.ts
@@ -228,6 +228,7 @@ export class Runtime {
       effect("nextState", targetState, () => {
         this.context.history.pop()
         this.context.history.push(targetState)
+        this.#contextDidChange()
       }),
 
       // Add a log effect.
@@ -249,12 +250,6 @@ export class Runtime {
 
       // Run enter on next state
       enter(),
-
-      // Notify listeners of change
-      // effect("contextChange", undefined, () => {
-      //   // Only state changes (and updates) can change context
-      //   this.onContextChange_()
-      // }),
     ]
 
     // Run exit on prior state first