[Workflows/vitest-pool-workers] Add a test handler to get the result of a workflow instance (#11648)

* [Workflows/vitest-pool-workers] Add Workflows result testing utils

* [Workflows/vitest-pool-workers] Add a test handler to get the result of a workflow instance

Author: pombosilva

.changeset/wicked-cities-attend.md

@@ -0,0 +1,28 @@
+---
+"@cloudflare/vitest-pool-workers": minor
+"@cloudflare/workflows-shared": minor
+"miniflare": minor
+---
+
+Add Workflows test handlers in vitest-pool-workers to get the Workflow instance output and error:
+
+- `getOutput()`: Returns the output of the successfully completed Workflow instance.
+- `getError()`: Returns the error information of the errored Workflow instance.
+
+Example:
+
+```ts
+// First wait for the workflow instance to complete:
+await expect(
+	instance.waitForStatus({ status: "complete" })
+).resolves.not.toThrow();
+
+// Then, get its output
+const output = await instance.getOutput();
+
+// Or for errored workflow instances, get their error:
+await expect(
+	instance.waitForStatus({ status: "errored" })
+).resolves.not.toThrow();
+const error = await instance.getError();
+```

fixtures/vitest-pool-workers-examples/workflows/test/integration.test.ts

@@ -26,6 +26,9 @@ it("workflow should be able to reach the end and be successful", async () => {
 	);
 	await expect(instance.waitForStatus(STATUS_COMPLETE)).resolves.not.toThrow();
 
+	const output = await instance.getOutput();
+	expect(output).toEqual({ status: "auto_approved" });
+
 	// DISPOSE: ensured by `await using`
 });
 
@@ -52,6 +55,9 @@ it("workflow batch should be able to reach the end and be successful", async ()
 			await expect(
 				instance.waitForStatus(STATUS_COMPLETE)
 			).resolves.not.toThrow();
+
+			const output = await instance.getOutput();
+			expect(output).toEqual({ status: "auto_approved" });
 		}
 	} finally {
 		// DISPOSE:

fixtures/vitest-pool-workers-examples/workflows/test/unit.test.ts

@@ -34,6 +34,9 @@ it("should mock a non-violation score and complete", async () => {
 
 	await expect(instance.waitForStatus(STATUS_COMPLETE)).resolves.not.toThrow();
 
+	const output = await instance.getOutput();
+	expect(output).toEqual({ status: "auto_approved" });
+
 	// DISPOSE: ensured by `await using`
 });
 
@@ -70,6 +73,9 @@ it("should mock the violation score calculation to fail 2 times and then complet
 		await expect(
 			instance.waitForStatus(STATUS_COMPLETE)
 		).resolves.not.toThrow();
+
+		const output = await instance.getOutput();
+		expect(output).toEqual({ status: "auto_approved" });
 	} finally {
 		// DISPOSE:
 		// Workflow introspector should be disposed the end of each test, if no `await using` dyntax is used
@@ -102,6 +108,9 @@ it("should mock a violation score and complete", async () => {
 	).toEqual({ status: "auto_rejected" });
 
 	await expect(instance.waitForStatus(STATUS_COMPLETE)).resolves.not.toThrow();
+
+	const output = await instance.getOutput();
+	expect(output).toEqual({ status: "auto_rejected" });
 });
 
 it("should be reviewed, accepted and complete", async () => {
@@ -132,6 +141,9 @@ it("should be reviewed, accepted and complete", async () => {
 	).toEqual({ status: "moderated", decision: "approve" });
 
 	await expect(instance.waitForStatus(STATUS_COMPLETE)).resolves.not.toThrow();
+
+	const output = await instance.getOutput();
+	expect(output).toEqual({ decision: "approve", status: "moderated" });
 });
 
 it("should force human review to timeout and error", async () => {
@@ -156,4 +168,8 @@ it("should force human review to timeout and error", async () => {
 	);
 
 	await expect(instance.waitForStatus(STATUS_ERROR)).resolves.not.toThrow();
+
+	const error = await instance.getError();
+	expect(error.name).toEqual("Error");
+	expect(error.message).toContain("Execution timed out");
 });

packages/miniflare/src/workers/workflows/wrapped-binding.worker.ts

@@ -54,6 +54,13 @@ class WorkflowImpl implements Workflow {
 	async unsafeWaitForStatus(instanceId: string, status: string): Promise<void> {
 		return await this.binding.unsafeWaitForStatus(instanceId, status);
 	}
+
+	public async unsafeGetOutputOrError(
+		instanceId: string,
+		isOutput: boolean
+	): Promise<unknown> {
+		return this.binding.unsafeGetOutputOrError(instanceId, isOutput);
+	}
 }
 
 class InstanceImpl implements WorkflowInstance {

packages/vitest-pool-workers/src/worker/workflows.ts

@@ -103,6 +103,17 @@ class WorkflowInstanceIntrospectorHandle
 		await this.#workflow.unsafeWaitForStatus(this.#instanceId, status);
 	}
 
+	async getOutput(): Promise<unknown> {
+		return await this.#workflow.unsafeGetOutputOrError(this.#instanceId, true);
+	}
+
+	async getError(): Promise<{ name: string; message: string }> {
+		return (await this.#workflow.unsafeGetOutputOrError(
+			this.#instanceId,
+			false
+		)) as { name: string; message: string };
+	}
+
 	async dispose(): Promise<void> {
 		await this.#workflow.unsafeAbort(this.#instanceId, "Instance dispose");
 	}

packages/vitest-pool-workers/types/cloudflare-test.d.ts

@@ -165,6 +165,9 @@ declare module "cloudflare:test" {
 	 *   // 3. ASSERTION
 	 *   await instance.waitForStatus("complete");
 	 *
+	 * 	 const output = await instance.getOutput();
+	 *   expect(output).toEqual({ success: true });
+	 *
 	 *   // 4. DISPOSE is implicit and automatic here.
 	 * });
 	 */
@@ -211,6 +214,49 @@ declare module "cloudflare:test" {
 		 */
 		waitForStatus(status: InstanceStatus["status"]): Promise<void>;
 
+		/**
+		 * Retrieves the output value returned by the Workflow instance upon successful completion.
+		 *
+		 * This method should only be called after the Workflow instance has completed successfully.
+		 * It's recommended to use `waitForStatus("complete")` before calling this method.
+		 *
+		 * @example
+		 * ```ts
+		 * it('my workflow test', async () => {
+		 *   await using instance = await introspectWorkflowInstance(env.MY_WORKFLOW, "123");
+		 *   await env.MY_WORKFLOW.create({ id: "123" });
+		 *
+		 *   await instance.waitForStatus("complete");
+		 *
+		 *   const output = await instance.getOutput();
+		 *   expect(output).toEqual({ success: true });
+		 * });
+		 * ```
+		 */
+		getOutput(): Promise<unknown>;
+
+		/**
+		 * Retrieves the error information from a failed Workflow instance.
+		 *
+		 * This method should only be called after the Workflow instance has failed.
+		 * It's recommended to use `waitForStatus("errored")` before calling this method.
+		 *
+		 * @example
+		 * ```ts
+		 * it('my workflow test', async () => {
+		 * 	 await using instance = await introspectWorkflowInstance(env.MY_WORKFLOW, "123");
+		 *   await env.MY_WORKFLOW.create({ id: "123" });
+		 *
+		 *   await instance.waitForStatus("errored");
+		 *
+		 *   const error = await instance.getError();
+		 *   expect(error.name).toBe("Error");
+		 *   expect(error.message).toContain("some error");
+		 * });
+		 * ```
+		 */
+		getError(): Promise<{ name: string; message: string }>;
+
 		/**
 		 * Disposes the Workflow instance introspector.
 		 *
@@ -433,6 +479,9 @@ declare module "cloudflare:test" {
 	 *   const instances = introspector.get();
 	 *   for(const instance of instances) {
 	 *     await instance.waitForStatus("complete");
+	 *
+	 *     const output = await instance.getOutput();
+	 *     expect(output).toEqual({ success: true });
 	 *   }
 	 *
 	 * // 4. DISPOSE is implicit and automatic here.

packages/workflows-shared/src/binding.ts

@@ -135,6 +135,15 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> {
 		const stub = this.env.ENGINE.get(stubId);
 		return await stub.waitForStatus(status);
 	}
+
+	public async unsafeGetOutputOrError(
+		instanceId: string,
+		isOutput: boolean
+	): Promise<unknown> {
+		const stubId = this.env.ENGINE.idFromName(instanceId);
+		const stub = this.env.ENGINE.get(stubId);
+		return await stub.getOutputOrError(isOutput);
+	}
 }
 
 export class WorkflowHandle extends RpcTarget implements WorkflowInstance {
@@ -167,17 +176,13 @@ export class WorkflowHandle extends RpcTarget implements WorkflowInstance {
 	public async status(): Promise<
 		InstanceStatus & { __LOCAL_DEV_STEP_OUTPUTS: unknown[] }
 	> {
-		const status = await this.stub.getStatus(0, this.id);
+		const status = await this.stub.getStatus();
 
 		// NOTE(lduarte): for some reason, sync functions over RPC are typed as never instead of Promise<EngineLogs>
 		using logs = await (this.stub.readLogs() as unknown as Promise<
 			EngineLogs & Disposable
 		>);
 
-		const workflowSuccessEvent = logs.logs
-			.filter((log) => log.event === InstanceEvent.WORKFLOW_SUCCESS)
-			.at(0);
-
 		const filteredLogs = logs.logs.filter(
 			(log) =>
 				log.event === InstanceEvent.STEP_SUCCESS ||
@@ -191,15 +196,19 @@ export class WorkflowHandle extends RpcTarget implements WorkflowInstance {
 		);
 
 		const workflowOutput =
-			workflowSuccessEvent !== undefined
-				? workflowSuccessEvent.metadata.result
-				: null;
+			logs.logs.find((log) => log.event === InstanceEvent.WORKFLOW_SUCCESS)
+				?.metadata.result ?? null;
+
+		const workflowError = logs.logs.find(
+			(log) => log.event === InstanceEvent.WORKFLOW_FAILURE
+		)?.metadata.error;
 
 		return {
 			status: instanceStatusName(status),
 			__LOCAL_DEV_STEP_OUTPUTS: stepOutputs,
 			output: workflowOutput,
-		}; // output, error
+			error: workflowError,
+		};
 	}
 
 	public async sendEvent(args: {

packages/workflows-shared/src/engine.ts

@@ -60,8 +60,9 @@ export type Log = {
 	group: string | null;
 	target: string | null;
 	metadata: {
-		result: unknown;
-		payload: unknown;
+		result: string;
+		payload: string;
+		error: { name: string; message: string };
 	};
 };
 
@@ -167,10 +168,29 @@ export class Engine extends DurableObject<Env> {
 		};
 	}
 
-	async getStatus(
-		_accountId: number,
-		_instanceId: string
-	): Promise<InstanceStatus> {
+	readLogsFromEvent(eventType: InstanceEvent): EngineLogs {
+		const logs = [
+			...this.ctx.storage.sql.exec<{
+				event: InstanceEvent;
+				groupKey: string | null;
+				target: string | null;
+				metadata: string;
+			}>(
+				"SELECT event, groupKey, target, metadata FROM states WHERE event = ?",
+				eventType
+			),
+		];
+
+		return {
+			logs: logs.map((log) => ({
+				...log,
+				metadata: JSON.parse(log.metadata),
+				group: log.groupKey,
+			})),
+		};
+	}
+
+	async getStatus(): Promise<InstanceStatus> {
 		if (this.accountId === undefined) {
 			// Engine could have restarted, so we try to restore from its state
 			const metadata =
@@ -353,6 +373,34 @@ export class Engine extends DurableObject<Env> {
 		}
 	}
 
+	async getOutputOrError(isOutput: boolean): Promise<unknown> {
+		const status = await this.getStatus();
+
+		if (isOutput) {
+			if (status !== InstanceStatus.Complete) {
+				throw new Error(
+					`Cannot retrieve output: Workflow instance is in status "${instanceStatusName(status)}" but must be "complete" to have an output available`
+				);
+			}
+			const logs = this.readLogsFromEvent(InstanceEvent.WORKFLOW_SUCCESS).logs;
+			return logs.at(0)?.metadata.result;
+		} else {
+			if (status !== InstanceStatus.Errored) {
+				throw new Error(
+					`Cannot retrieve error: Workflow instance is in status "${instanceStatusName(status)}" but must be "errored" to have error information available`
+				);
+			}
+			const logs = this.readLogsFromEvent(InstanceEvent.WORKFLOW_FAILURE).logs;
+			const log = logs.at(0);
+			if (!log?.metadata.error) {
+				throw new Error(
+					"Cannot retrieve error: No workflow instance failure log found"
+				);
+			}
+			return log.metadata.error;
+		}
+	}
+
 	async abort(_reason: string) {
 		// TODO: Maybe don't actually kill but instead check a flag and return early if true
 	}
@@ -497,7 +545,7 @@ export class Engine extends DurableObject<Env> {
 		this.instanceId = instance.id;
 		this.workflowName = workflow.name;
 
-		const status = await this.getStatus(accountId, instance.id);
+		const status = await this.getStatus();
 		if (
 			[
 				InstanceStatus.Errored, // TODO (WOR-85): Remove this once upgrade story is done

packages/workflows-shared/tests/engine.test.ts

@@ -252,7 +252,7 @@ describe("Engine", () => {
 		const restartedStub = env.ENGINE.get(engineId);
 
 		const status = await runInDurableObject(restartedStub, (engine) => {
-			return engine.getStatus(accountId, instanceId);
+			return engine.getStatus();
 		});
 
 		expect(status).toBe(InstanceStatus.Running);