From 5a7979153e0a74084b35fccdc62ed8896fa5138e Mon Sep 17 00:00:00 2001
From: Thomas Reynolds <me@tdreyno.com>
Date: Fri, 5 May 2023 12:54:39 -0700
Subject: [PATCH] fix: allow async wait state handlers

---
 src/__tests__/waitState.spec.ts | 49 +++++++++++++++++++++++++++++++++
 src/state.ts                    |  4 +--
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/src/__tests__/waitState.spec.ts b/src/__tests__/waitState.spec.ts
index 87800cc5..40f9f483 100644
--- a/src/__tests__/waitState.spec.ts
+++ b/src/__tests__/waitState.spec.ts
@@ -120,4 +120,53 @@ describe("waitState", () => {
     expect(isState(runtime.currentState(), TimedOutState)).toBeTruthy()
     expect((runtime.currentState().data as D).count).toBe(INITIAL_COUNT)
   })
+
+  it("should accept async handlers", async () => {
+    const BeforeThing = state<Enter, D>(
+      {
+        Enter: data => WaitForThing([data, RETURN_COUNT]),
+      },
+      { name: "BeforeThing" },
+    )
+
+    const TimedOutState = state<Enter, D>(
+      {
+        Enter: noop,
+      },
+      { name: "TimedOutState" },
+    )
+
+    const WaitForThing = waitState(
+      fetchThing,
+      thingFetched,
+      async (data: D, payload) => {
+        return AfterThing({ ...data, count: data.count + payload })
+      },
+      {
+        name: "WaitForThing",
+        timeout: 1000,
+        onTimeout: async data => {
+          return TimedOutState(data)
+        },
+      },
+    )
+    const context = createInitialContext([
+      BeforeThing({
+        count: INITIAL_COUNT,
+      }),
+    ])
+
+    const runtime = createRuntime(context, {}, { fetchThing })
+
+    expect(isState(runtime.currentState(), BeforeThing)).toBeTruthy()
+
+    await runtime.run(enter())
+
+    expect(isState(runtime.currentState(), WaitForThing)).toBeTruthy()
+
+    await timeout(2000)
+
+    expect(isState(runtime.currentState(), TimedOutState)).toBeTruthy()
+    expect((runtime.currentState().data as D).count).toBe(INITIAL_COUNT)
+  })
 })
diff --git a/src/state.ts b/src/state.ts
index 08e50ad1..ceb15d96 100644
--- a/src/state.ts
+++ b/src/state.ts
@@ -330,11 +330,11 @@ export const waitState = <
 >(
   requestAction: ReqAC,
   responseActionCreator: RespAC,
-  transition: (data: Data, payload: RespA["payload"]) => StateReturn,
+  transition: (data: Data, payload: RespA["payload"]) => HandlerReturn,
   options?: {
     name?: string
     timeout?: number
-    onTimeout?: (data: Data) => StateReturn
+    onTimeout?: (data: Data) => HandlerReturn
   },
 ) => {
   const name = options?.name