fix: add get waiters method and try to parse notify/wait response by defualt

Author: CahidArda

src/client/index.ts

@@ -1,6 +1,6 @@
-import { NotifyResponse } from "../types";
+import { NotifyResponse, Waiter } from "../types";
 import { Client as QStashClient } from "@upstash/qstash";
-import { makeNotifyRequest } from "./utils";
+import { makeGetWaitersRequest, makeNotifyRequest } from "./utils";
 
 type ClientConfig = ConstructorParameters<typeof QStashClient>[0];
 
@@ -42,6 +42,15 @@ export class Client {
     eventId: string;
     eventData?: unknown;
   }): Promise<NotifyResponse[]> {
-    return await makeNotifyRequest(this.client.http, eventId, eventData)
+    return await makeNotifyRequest(this.client.http, eventId, eventData);
+  }
+
+  /**
+   * Check waiters of an event
+   *
+   * @param eventId event id to check
+   */
+  public async getWaiters({ eventId }: { eventId: string }): Promise<Waiter[]> {
+    return await makeGetWaitersRequest(this.client.http, eventId);
   }
 }

src/client/utils.ts

@@ -1,5 +1,5 @@
 import { Client } from "@upstash/qstash";
-import { NotifyResponse } from "../types";
+import { NotifyResponse, Waiter } from "../types";
 
 export const makeNotifyRequest = async (
   requester: Client["http"],
@@ -13,4 +13,15 @@ export const makeNotifyRequest = async (
   })) as NotifyResponse[];
 
   return result;
-}
+};
+
+export const makeGetWaitersRequest = async (
+  requester: Client["http"],
+  eventId: string
+): Promise<Waiter[]> => {
+  const result = (await requester.request({
+    path: ["v2", "waiters", eventId],
+    method: "GET",
+  })) as Waiter[];
+  return result;
+};

src/context/context.ts

@@ -341,24 +341,33 @@ export class WorkflowContext<TInitialPayload = unknown> {
       )
     );
 
-    return result;
+    try {
+      return {
+        ...result,
+        eventData: JSON.parse(result.eventData as string),
+      };
+    } catch {
+      return result;
+    }
   }
 
   public async notify(
     stepName: string,
     eventId: string,
-    eventData: string
+    eventData: unknown
   ): Promise<NotifyStepResponse> {
     const result = await this.addStep(
-      new LazyNotifyStep(
-        stepName,
-        eventId,
-        eventData,
-        this.qstashClient.http
-      )
+      new LazyNotifyStep(stepName, eventId, eventData, this.qstashClient.http)
     );
 
-    return result;
+    try {
+      return {
+        ...result,
+        eventData: JSON.parse(result.eventData as string),
+      };
+    } catch {
+      return result;
+    }
   }
 
   /**

src/context/steps.test.ts

@@ -186,11 +186,11 @@ describe("test steps", () => {
 
   describe("notify step", () => {
     const eventId = "my-event-id";
-    const eventData = {data: "my-event-data"};
+    const eventData = { data: "my-event-data" };
 
     // get client
     const token = nanoid();
-    const client = new Client({ baseUrl: MOCK_QSTASH_SERVER_URL, token })
+    const client = new Client({ baseUrl: MOCK_QSTASH_SERVER_URL, token });
 
     const step = new LazyNotifyStep(stepName, eventId, eventData, client.http);
 
@@ -209,7 +209,7 @@ describe("test steps", () => {
     });
 
     test("should create result step", async () => {
-      let called = false
+      let called = false;
       const notifyResponse: NotifyResponse[] = [
         {
           error: "no-error",
@@ -224,27 +224,27 @@ describe("test steps", () => {
               "my-header": ["value"],
             },
             timeoutUrl: "url",
-            url: "url"
-          }
-        }
-      ]
+            url: "url",
+          },
+        },
+      ];
       const stepResponse: NotifyStepResponse = {
         eventId,
         eventData,
-        notifyResponse
-      }
+        notifyResponse,
+      };
 
       await mockQStashServer({
         execute: async () => {
-          const result = await step.getResultStep(4, stepId)
+          const result = await step.getResultStep(4, stepId);
           expect(result).toEqual({
             concurrent: 4,
             stepId,
             out: stepResponse,
             stepName,
             stepType: "Notify",
           });
-          called = true
+          called = true;
         },
         responseFields: {
           status: 200,
@@ -256,10 +256,9 @@ describe("test steps", () => {
           token,
           body: eventData,
         },
-      })
+      });
 
-      expect(called).toBeTrue()
+      expect(called).toBeTrue();
     });
-
-  })
+  });
 });

src/context/steps.ts

@@ -227,26 +227,17 @@ export class LazyWaitForEventStep extends BaseLazyStep<WaitStepResponse> {
 }
 
 export class LazyNotifyStep extends LazyFunctionStep<NotifyStepResponse> {
-  stepType: StepType = "Notify"
+  stepType: StepType = "Notify";
 
-  constructor(
-    stepName: string,
-    eventId: string,
-    eventData: unknown,
-    requester: Client["http"]
-  ) {
+  constructor(stepName: string, eventId: string, eventData: unknown, requester: Client["http"]) {
     super(stepName, async () => {
-      const notifyResponse = await makeNotifyRequest(
-        requester,
-        eventId,
-        eventData
-      )
+      const notifyResponse = await makeNotifyRequest(requester, eventId, eventData);
 
       return {
         eventId,
         eventData,
-        notifyResponse
-      }
-    })
+        notifyResponse,
+      };
+    });
   }
 }

src/integration.test.ts

@@ -55,10 +55,10 @@ import type { RouteFunction, WaitStepResponse, WorkflowServeOptions } from "./ty
 import type { NextRequest } from "next/server";
 import { Client } from "./client";
 import { nanoid } from "./utils";
+import { makeGetWaitersRequest } from "./client/utils";
 
 const WORKFLOW_PORT = "3000";
 const THIRD_PARTY_PORT = "3001";
-const LOCAL_WORKFLOW_URL = `http://localhost:${WORKFLOW_PORT}`;
 const LOCAL_THIRD_PARTY_URL = `http://localhost:${THIRD_PARTY_PORT}`;
 
 const someWork = (input: string) => {
@@ -109,21 +109,23 @@ const testEndpoint = async <TInitialPayload = unknown>({
   finishState,
   failureFunction,
   retries,
+  port = WORKFLOW_PORT,
 }: {
-  finalCount: number;
+  finalCount?: number;
   waitFor: number;
   initialPayload: TInitialPayload;
   routeFunction: RouteFunction<TInitialPayload>;
   finishState: FinishState;
   failureFunction?: WorkflowServeOptions["failureFunction"];
   retries?: number;
+  port?: string;
 }) => {
   let counter = 0;
 
   const endpoint = workflowServe<TInitialPayload>(routeFunction, {
     qstashClient,
-    url: LOCAL_WORKFLOW_URL,
-    verbose: true,
+    url: `http://localhost:${port}`,
+    // verbose: true,
     failureFunction,
     retries,
   });
@@ -133,7 +135,7 @@ const testEndpoint = async <TInitialPayload = unknown>({
       counter += 1;
       return await endpoint(request as NextRequest);
     },
-    port: WORKFLOW_PORT,
+    port: port,
   });
 
   await qstashClient.publishJSON({
@@ -142,15 +144,17 @@ const testEndpoint = async <TInitialPayload = unknown>({
     headers: {
       Authentication: "Bearer secretPassword",
     },
-    url: `http://localhost:${WORKFLOW_PORT}`,
+    url: `http://localhost:${port}`,
   });
 
   await new Promise((resolve) => setTimeout(resolve, waitFor));
 
   server.stop();
 
   finishState.check();
-  expect(counter).toBe(finalCount);
+  if (finalCount) {
+    expect(counter).toBe(finalCount);
+  }
 };
 
 describe.skip("live serve tests", () => {
@@ -590,8 +594,9 @@ describe.skip("live serve tests", () => {
             context.waitForEvent("wait-event-step", eventId, 3),
           ]);
           expect(runResponse).toBe(runResult);
-          expect(waitResponse.eventData).toBe(expectedWaitResponse.eventData);
           expect(waitResponse.timeout).toBe(expectedWaitResponse.timeout);
+          expect(waitResponse.eventData).toEqual(expectedWaitResponse.eventData);
+          expect(typeof waitResponse.eventData).toBe(typeof expectedWaitResponse.eventData);
           finishState.finish();
         },
       });
@@ -651,5 +656,64 @@ describe.skip("live serve tests", () => {
       },
       { timeout: 17_000 }
     );
+
+    describe("should notify from inside a function", () => {
+      const testNotifyWithContext = async (payload: unknown) => {
+        const eventId = `my-event-id-${nanoid()}`;
+
+        const waitingEndpoint = testWaitEndpoint(
+          {
+            eventData: payload,
+            timeout: false,
+          },
+          eventId
+        );
+
+        const finishState = new FinishState();
+        const notifyingEndpoint = testEndpoint({
+          finishState,
+          initialPayload: undefined,
+          waitFor: 15000,
+          port: "3002",
+          routeFunction: async (context) => {
+            // wait to avoid notifying the first waitForEvent
+            await context.sleep("sleep for first timeout", 3);
+
+            while (true) {
+              const waiterExists = await context.run("check waiters", async () => {
+                const waiters = await makeGetWaitersRequest(context.qstashClient.http, eventId);
+
+                return Boolean(waiters);
+              });
+
+              if (waiterExists) {
+                break;
+              }
+            }
+            const { notifyResponse } = await context.notify("notify-step", eventId, payload);
+            expect(notifyResponse.length).toBeTruthy();
+            finishState.finish();
+          },
+        });
+
+        await Promise.all([waitingEndpoint, notifyingEndpoint]);
+      };
+
+      test(
+        "should handle string event data",
+        async () => {
+          await testNotifyWithContext("event-data");
+        },
+        { timeout: 170000 }
+      );
+
+      test(
+        "should handle object event data",
+        async () => {
+          await testNotifyWithContext({ event: "data" });
+        },
+        { timeout: 170000 }
+      );
+    });
   });
 });

src/types.ts

@@ -275,13 +275,13 @@ export type NotifyStepResponse = {
   /**
    * notified event id
    */
-  eventId: string,
+  eventId: string;
   /**
    * event data sent with notify
    */
-  eventData: unknown,
+  eventData: unknown;
   /**
    * response from notify
    */
-  notifyResponse: NotifyResponse[]
-}
+  notifyResponse: NotifyResponse[];
+};